package client

import (
	"fmt"
	"log"

	"github.com/duffy/usb-server/internal/crypto"
)

// tunnelCodec seals and opens tunnel payloads.
//
// It sits above the transport so that a tunnel is protected the same way
// whether its frames travel directly or through the relay. A nil codec passes
// data through unchanged, which is what a client configured with only a group
// hash — and therefore unable to derive the key — falls back to.
type tunnelCodec struct {
	sealer *crypto.Sealer
	opener *crypto.Opener
}

// newTunnelCodec builds a codec for one end of a tunnel.
// send is the direction this end transmits in; it receives on the other.
func newTunnelCodec(secret *crypto.TunnelSecret, tunnelID string, send crypto.Direction) (*tunnelCodec, error) {
	if secret == nil {
		return nil, nil // unencrypted, by configuration
	}

	key, err := secret.TunnelKey(tunnelID)
	if err != nil {
		return nil, err
	}

	recv := crypto.DirShareToUse
	if send == crypto.DirShareToUse {
		recv = crypto.DirUseToShare
	}

	sealer, err := crypto.NewSealer(key, send)
	if err != nil {
		return nil, err
	}
	opener, err := crypto.NewOpener(key, recv)
	if err != nil {
		return nil, err
	}

	return &tunnelCodec{sealer: sealer, opener: opener}, nil
}

// encode prepares a payload for transmission.
func (c *tunnelCodec) encode(payload []byte) ([]byte, error) {
	if c == nil {
		return payload, nil
	}
	return c.sealer.Seal(payload)
}

// decode recovers a received payload.
func (c *tunnelCodec) decode(frame []byte) ([]byte, error) {
	if c == nil {
		return frame, nil
	}
	return c.opener.Open(frame)
}

// encrypted reports whether this codec actually protects anything.
func (c *tunnelCodec) encrypted() bool { return c != nil }

// tunnelSender delivers one encoded frame to the peer.
type tunnelSender func(frame []byte) error

// relaySender routes frames through the relay, tagged with the tunnel ID.
func relaySender(c *Client, tunnelID string) tunnelSender {
	return func(frame []byte) error {
		return c.SendTunnelData(tunnelID, frame)
	}
}

// directSender routes frames over an established direct connection.
func directSender(conn *directConn) tunnelSender {
	return func(frame []byte) error {
		return conn.WriteFrame(frame)
	}
}

// send encodes a payload and hands it to the transport.
func send(codec *tunnelCodec, sender tunnelSender, payload []byte) error {
	frame, err := codec.encode(payload)
	if err != nil {
		return fmt.Errorf("encoding tunnel frame: %w", err)
	}
	return sender(frame)
}

// receiveLoop reads frames from a direct connection, decodes them and hands
// each payload to deliver. It returns when the connection ends, when the
// tunnel is torn down, or on the first frame that fails to authenticate.
func receiveLoop(conn *directConn, codec *tunnelCodec, deliver func([]byte) error, done <-chan struct{}, label string) {
	for {
		select {
		case <-done:
			return
		default:
		}

		frame, err := conn.ReadFrame()
		if err != nil {
			select {
			case <-done:
			default:
				log.Printf("[direct] %s: read ended: %v", label, err)
			}
			return
		}

		payload, err := codec.decode(frame)
		if err != nil {
			// A frame that fails to authenticate means the stream is either
			// corrupt or being tampered with. Either way this tunnel cannot
			// be trusted to carry USB traffic any further.
			log.Printf("[direct] %s: dropping connection: %v", label, err)
			return
		}

		if err := deliver(payload); err != nil {
			log.Printf("[direct] %s: delivery failed: %v", label, err)
			return
		}
	}
}