/*
Copyright 2022 CAcert Inc.
SPDX-License-Identifier: Apache-2.0

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package protocol handles the protocol message marshaling and unmarshalling.
package protocol

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"time"

	"github.com/justincpresley/go-cobs"
	"github.com/sirupsen/logrus"

	"git.cacert.org/cacert-gosigner/pkg/messages"
)

const CobsDelimiter = 0x00

type Command struct {
	Announce *messages.CommandAnnounce
	Command  interface{}
}

func (c *Command) String() string {
	return fmt.Sprintf("Cmd[announce={%s}, data={%s}]", c.Announce, c.Command)
}

type Response struct {
	Announce *messages.ResponseAnnounce
	Response interface{}
}

func (r *Response) String() string {
	return fmt.Sprintf("Rsp[announce={%s}, data={%s}]", r.Announce, r.Response)
}

// ServerHandler is responsible for parsing incoming frames and calling commands
type ServerHandler interface {
	// CommandAnnounce handles the initial announcement of a command.
	CommandAnnounce(context.Context, <-chan []byte) (*Command, error)
	// CommandData handles the command data.
	CommandData(context.Context, <-chan []byte, *Command) error
	// HandleCommand executes the command, generating a response.
	HandleCommand(context.Context, *Command) (*Response, error)
	// Respond generates the response for a command.
	Respond(context.Context, *Response, chan<- []byte) error
}

type ClientHandler interface {
	Send(context.Context, *Command, chan<- []byte) error
	ResponseAnnounce(context.Context, <-chan []byte) (*Response, error)
	ResponseData(context.Context, <-chan []byte, *Response) error
	HandleResponse(context.Context, *Response) error
}

var (
	errCommandExpected         = errors.New("command must not be nil")
	errCommandAnnounceExpected = errors.New("command must have an announcement")
	errCommandDataExpected     = errors.New("command must have data")
	errResponseExpected        = errors.New("response must not be nil")

	ErrResponseAnnounceTimeoutExpired = errors.New("response announce timeout expired")
	ErrResponseDataTimeoutExpired     = errors.New("response data timeout expired")
)

type protocolState int8

const (
	cmdAnnounce protocolState = iota
	cmdData
	handleCommand
	respond
	respAnnounce
	respData
	handleResponse
)

var protocolStateNames = map[protocolState]string{
	cmdAnnounce:    "CMD ANNOUNCE",
	cmdData:        "CMD DATA",
	handleCommand:  "HANDLE CMD",
	respond:        "RESPOND",
	respAnnounce:   "RESP ANNOUNCE",
	respData:       "RESP DATA",
	handleResponse: "HANDLE RESP",
}

func (p protocolState) String() string {
	if name, ok := protocolStateNames[p]; ok {
		return name
	}

	return fmt.Sprintf("unknown %d", p)
}

type ServerProtocol struct {
	handler ServerHandler
	in      <-chan []byte
	out     chan<- []byte
	logger  *logrus.Logger
	state   protocolState
}

func (p *ServerProtocol) Handle(ctx context.Context) error {
	var (
		command  *Command
		response *Response
		err      error
	)

	for {
		select {
		case <-ctx.Done():
			close(p.out)

			return nil
		default:
			p.logger.Debugf("handling protocol state %s", p.state)

			switch p.state {
			case cmdAnnounce:
				command = p.commandAnnounce(ctx)
			case cmdData:
				err = p.commandData(ctx, command)
				if err != nil {
					return err
				}
			case handleCommand:
				response, err = p.handleCommand(ctx, command)
				if err != nil {
					return err
				}
			case respond:
				err = p.respond(ctx, response)
				if err != nil {
					return err
				}
			default:
				return fmt.Errorf("unknown protocol state %s", p.state)
			}
		}
	}
}

func (p *ServerProtocol) commandAnnounce(ctx context.Context) *Command {
	command, err := p.handler.CommandAnnounce(ctx, p.in)
	if err != nil {
		p.logger.WithError(err).Error("could not handle command announce")

		return nil
	}

	p.state = cmdData

	return command
}

func (p *ServerProtocol) commandData(ctx context.Context, command *Command) error {
	if command == nil || command.Announce == nil {
		return errCommandAnnounceExpected
	}

	err := p.handler.CommandData(ctx, p.in, command)
	if err != nil {
		p.logger.WithError(err).Error("could not handle command data")

		p.state = cmdAnnounce

		return nil
	}

	p.state = handleCommand

	return nil
}

func (p *ServerProtocol) handleCommand(ctx context.Context, command *Command) (*Response, error) {
	if command == nil || command.Announce == nil || command.Command == nil {
		return nil, errCommandDataExpected
	}

	response, err := p.handler.HandleCommand(ctx, command)
	if err != nil {
		p.logger.WithError(err).Error("could not handle command")

		p.state = cmdAnnounce

		return nil, nil
	}

	p.state = respond

	return response, nil
}

func (p *ServerProtocol) respond(ctx context.Context, response *Response) error {
	if response == nil {
		return errResponseExpected
	}

	err := p.handler.Respond(ctx, response, p.out)
	if err != nil {
		p.logger.WithError(err).Error("could not respond")

		p.state = cmdAnnounce

		return nil
	}

	p.state = cmdAnnounce

	return nil
}

func NewServer(handler ServerHandler, in <-chan []byte, out chan<- []byte, logger *logrus.Logger) *ServerProtocol {
	return &ServerProtocol{
		handler: handler,
		in:      in,
		out:     out,
		logger:  logger,
		state:   cmdAnnounce,
	}
}

type ClientProtocol struct {
	handler  ClientHandler
	commands <-chan *Command
	in       <-chan []byte
	out      chan<- []byte
	logger   *logrus.Logger
	state    protocolState
}

func (p *ClientProtocol) Handle(ctx context.Context) error {
	var (
		response *Response
		err      error
	)

	for {
		select {
		case <-ctx.Done():
			return nil
		default:
			p.logger.Debugf("handling protocol state %s", p.state)

			switch p.state {
			case cmdAnnounce:
				err = p.cmdAnnounce(ctx)
				if err != nil {
					return err
				}
			case respAnnounce:
				response = p.respAnnounce(ctx)
			case respData:
				err = p.respData(ctx, response)
				if err != nil {
					return err
				}
			case handleResponse:
				err = p.handleResponse(ctx, response)
				if err != nil {
					return err
				}
			default:
				return fmt.Errorf("unknown protocol state %s", p.state)
			}
		}
	}
}

func (p *ClientProtocol) cmdAnnounce(ctx context.Context) error {
	select {
	case <-ctx.Done():
		return nil
	case command, ok := <-p.commands:
		if !ok {
			return errCommandExpected
		}

		err := p.handler.Send(ctx, command, p.out)
		if err != nil {
			p.logger.WithError(err).Error("could not send command announce")

			return nil
		}

		p.logger.WithFields(map[string]interface{}{
			"command":       command.Announce,
			"buffer length": len(p.commands),
		}).Trace("handled command")
	}

	p.state = respAnnounce

	return nil
}

func (p *ClientProtocol) respAnnounce(ctx context.Context) *Response {
	response, err := p.handler.ResponseAnnounce(ctx, p.in)
	if err != nil {
		p.logger.WithError(err).Error("could not handle response announce")

		p.state = cmdAnnounce

		return nil
	}

	p.state = respData

	return response
}

func (p *ClientProtocol) respData(ctx context.Context, response *Response) error {
	if response == nil || response.Announce == nil {
		return errResponseExpected
	}

	err := p.handler.ResponseData(ctx, p.in, response)
	if err != nil {
		p.logger.WithError(err).Error("could not handle response data")

		if errors.Is(err, ErrResponseDataTimeoutExpired) {
			p.state = cmdAnnounce
		} else {
			p.state = respAnnounce
		}

		return nil
	}

	p.state = handleResponse

	return nil
}

func (p *ClientProtocol) handleResponse(ctx context.Context, response *Response) error {
	if response == nil || response.Announce == nil || response.Response == nil {
		return errResponseExpected
	}

	err := p.handler.HandleResponse(ctx, response)
	if err != nil {
		p.logger.WithError(err).Error("could not handle response")

		p.state = cmdAnnounce

		return nil
	}

	p.state = cmdAnnounce

	return nil
}

func NewClient(
	handler ClientHandler,
	commands <-chan *Command,
	in <-chan []byte,
	out chan<- []byte,
	logger *logrus.Logger,
) *ClientProtocol {
	return &ClientProtocol{
		handler:  handler,
		commands: commands,
		in:       in,
		out:      out,
		logger:   logger,
		state:    cmdAnnounce,
	}
}

// Framer handles bytes on the wire by adding or removing framing information.
type Framer interface {
	// ReadFrames reads data frames and publishes unframed data to the channel.
	ReadFrames(context.Context, io.Reader, chan<- []byte) error
	// WriteFrames takes data from the channel and writes framed data to the writer.
	WriteFrames(context.Context, io.Writer, <-chan []byte) error
}

const bufferSize = 1024
const readInterval = 50 * time.Millisecond

type COBSFramer struct {
	logger  *logrus.Logger
	encoder cobs.Encoder
}

func NewCOBSFramer(logger *logrus.Logger) (*COBSFramer, error) {
	encoder, err := cobs.NewEncoder(cobs.Config{SpecialByte: CobsDelimiter, Delimiter: true, EndingSave: true})
	if err != nil {
		return nil, fmt.Errorf("could not setup encoder: %w", err)
	}

	return &COBSFramer{
		encoder: encoder,
		logger:  logger,
	}, nil
}

func (c *COBSFramer) ReadFrames(ctx context.Context, reader io.Reader, frameChan chan<- []byte) error {
	var (
		err              error
		raw, data, frame []byte
	)

	buffer := &bytes.Buffer{}

	for {
		select {
		case <-ctx.Done():
			return nil
		default:
			raw, err = c.readRaw(ctx, reader)
			if err != nil {
				close(frameChan)

				return err
			}

			if len(raw) == 0 {
				time.Sleep(readInterval)

				continue
			}

			c.logger.Tracef("read %d raw bytes", len(raw))

			buffer.Write(raw)

			for {
				data, err = buffer.ReadBytes(CobsDelimiter)
				if err != nil {
					if errors.Is(err, io.EOF) {
						buffer.Write(data)

						break
					}

					// this is a safety measure, buffer.ReadBytes should only return io.EOF
					return fmt.Errorf("could not read from buffer: %w", err)
				}

				if err = c.encoder.Verify(data); err != nil {
					c.logger.WithError(err).Warnf("skipping invalid frame of %d bytes", len(data))

					continue
				}

				frame = c.encoder.Decode(data)

				c.logger.Tracef("frame decoded to length %d", len(frame))

				frameChan <- frame
			}

			c.logger.Tracef("read buffer is now %d bytes long", buffer.Len())
		}
	}
}

func (c *COBSFramer) readRaw(ctx context.Context, reader io.Reader) ([]byte, error) {
	result := make(chan []byte)
	errChan := make(chan error)

	go func() {
		buf := make([]byte, bufferSize)

		count, err := reader.Read(buf)
		if err != nil {
			if !errors.Is(err, io.EOF) {
				errChan <- fmt.Errorf("could not read data: %w", err)
			}
		}

		result <- buf[:count]
	}()

	select {
	case <-ctx.Done():
		return nil, nil
	case raw := <-result:
		return raw, nil
	case err := <-errChan:
		return nil, err
	}
}

func (c *COBSFramer) WriteFrames(ctx context.Context, writer io.Writer, frameChan <-chan []byte) error {
	for {
		select {
		case <-ctx.Done():
			return nil
		case frame := <-frameChan:
			if frame == nil {
				c.logger.Debug("channel closed")

				return nil
			}

			encoded := c.encoder.Encode(frame)

			err := c.writeRaw(ctx, writer, encoded)
			if err != nil {
				return err
			}
		}
	}
}

func (c *COBSFramer) writeRaw(ctx context.Context, writer io.Writer, raw []byte) error {
	errChan := make(chan error)

	go func() {
		n, err := io.Copy(writer, bytes.NewReader(raw))
		if err != nil {
			errChan <- fmt.Errorf("could not write data: %w", err)

			return
		}

		c.logger.WithField("count", n).Tracef("wrote bytes")

		close(errChan)
	}()

	select {
	case <-ctx.Done():
		return nil
	case err := <-errChan:
		return err
	}
}