package client

import (
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

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

// Direct tunnel wire format.
//
// Handshake, sent by the connecting (use) side:
//
//	[4] magic "USBD"  [1] version  [16] tunnel ID  [32] peer token
//
// Reply, sent by the listening (share) side:
//
//	[4] magic "USBD"  [1] version  [1] status (0 = accepted)
//
// Everything after that is length-prefixed encrypted frames:
//
//	[4] length (big endian)  [length bytes] sealed frame
const (
	directMagic   = "USBD"
	directVersion = 1

	handshakeSize      = 4 + 1 + protocol.TunnelHeaderSize + 32
	handshakeReplySize = 4 + 1 + 1

	// directHandshakeTimeout bounds the handshake. A peer that reaches the
	// port but does not speak this protocol must not hold the slot.
	directHandshakeTimeout = 5 * time.Second

	// directDialTimeout bounds one connection attempt. Candidate addresses
	// are tried in parallel, so this is also how long the whole attempt takes
	// before falling back to the relay.
	directDialTimeout = 3 * time.Second

	// maxDirectFrame caps a single frame, so a hostile or corrupt length
	// prefix cannot make us allocate arbitrarily.
	maxDirectFrame = 2 << 20
)

// Handshake status codes.
const (
	directAccepted    = 0
	directUnknownTun  = 1
	directBadToken    = 2
	directWrongVerson = 3
)

// directConn carries length-prefixed frames over a plain TCP connection.
//
// It deliberately does no encryption of its own: tunnel frames are sealed one
// level up, by the tunnel's codec, so that relayed and direct tunnels get the
// same protection. Putting it here instead would leave the relay path in
// cleartext — the one path where a third party is actually in the middle.
type directConn struct {
	conn net.Conn

	writeMu sync.Mutex
}

func newDirectConn(conn net.Conn) *directConn {
	return &directConn{conn: conn}
}

// WriteFrame sends one length-prefixed frame.
func (d *directConn) WriteFrame(payload []byte) error {
	if len(payload) > maxDirectFrame {
		return fmt.Errorf("frame of %d bytes exceeds the %d byte limit", len(payload), maxDirectFrame)
	}

	buf := make([]byte, 4+len(payload))
	binary.BigEndian.PutUint32(buf, uint32(len(payload)))
	copy(buf[4:], payload)

	// TCP writes from several goroutines would interleave and corrupt the
	// framing, so sends are serialised.
	d.writeMu.Lock()
	defer d.writeMu.Unlock()

	if _, err := d.conn.Write(buf); err != nil {
		return fmt.Errorf("writing frame: %w", err)
	}
	return nil
}

// ReadFrame reads one length-prefixed frame.
func (d *directConn) ReadFrame() ([]byte, error) {
	var lenBuf [4]byte
	if _, err := io.ReadFull(d.conn, lenBuf[:]); err != nil {
		return nil, err
	}

	length := binary.BigEndian.Uint32(lenBuf[:])
	if length == 0 || length > maxDirectFrame {
		return nil, fmt.Errorf("frame length %d is out of range", length)
	}

	payload := make([]byte, length)
	if _, err := io.ReadFull(d.conn, payload); err != nil {
		return nil, err
	}
	return payload, nil
}

// RemoteAddr reports the peer address, for logging.
func (d *directConn) RemoteAddr() string { return d.conn.RemoteAddr().String() }

// Close closes the underlying connection.
func (d *directConn) Close() error { return d.conn.Close() }

// buildHandshake assembles the greeting the connecting side sends.
func buildHandshake(tunnelID, peerToken string) ([]byte, error) {
	tokenBytes, err := hex.DecodeString(peerToken)
	if err != nil || len(tokenBytes) != 32 {
		return nil, fmt.Errorf("invalid peer token")
	}
	if len(tunnelID) != protocol.TunnelHeaderSize {
		return nil, fmt.Errorf("tunnel ID is %d bytes, want %d", len(tunnelID), protocol.TunnelHeaderSize)
	}

	buf := make([]byte, 0, handshakeSize)
	buf = append(buf, directMagic...)
	buf = append(buf, directVersion)
	buf = append(buf, tunnelID...)
	buf = append(buf, tokenBytes...)
	return buf, nil
}

// parseHandshake validates the greeting and returns the requested tunnel ID
// and the presented token in hex form.
func parseHandshake(data []byte) (tunnelID, peerToken string, err error) {
	if len(data) != handshakeSize {
		return "", "", fmt.Errorf("handshake is %d bytes, want %d", len(data), handshakeSize)
	}
	if string(data[:4]) != directMagic {
		return "", "", fmt.Errorf("bad magic")
	}
	if data[4] != directVersion {
		return "", "", fmt.Errorf("unsupported version %d", data[4])
	}

	tunnelID = string(data[5 : 5+protocol.TunnelHeaderSize])
	peerToken = hex.EncodeToString(data[5+protocol.TunnelHeaderSize:])
	return tunnelID, peerToken, nil
}

func buildHandshakeReply(status byte) []byte {
	buf := make([]byte, 0, handshakeReplySize)
	buf = append(buf, directMagic...)
	buf = append(buf, directVersion)
	buf = append(buf, status)
	return buf
}

func parseHandshakeReply(data []byte) error {
	if len(data) != handshakeReplySize {
		return fmt.Errorf("reply is %d bytes, want %d", len(data), handshakeReplySize)
	}
	if string(data[:4]) != directMagic {
		return fmt.Errorf("bad magic in reply")
	}
	if data[4] != directVersion {
		return fmt.Errorf("peer speaks version %d, we speak %d", data[4], directVersion)
	}

	switch data[5] {
	case directAccepted:
		return nil
	case directUnknownTun:
		return fmt.Errorf("peer does not know this tunnel")
	case directBadToken:
		return fmt.Errorf("peer rejected our token")
	case directWrongVerson:
		return fmt.Errorf("peer rejected our version")
	default:
		return fmt.Errorf("peer rejected the connection (status %d)", data[5])
	}
}

// localEndpoints lists host:port addresses on this machine's own interfaces.
//
// Loopback is skipped — a peer on another machine cannot use it — but every
// other usable unicast address is offered, because which one is reachable
// depends on the network and only the attempt can tell.
func localEndpoints(port int) []string {
	if port == 0 {
		return nil
	}

	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return nil
	}

	var endpoints []string
	for _, addr := range addrs {
		ipNet, ok := addr.(*net.IPNet)
		if !ok {
			continue
		}
		ip := ipNet.IP
		if ip.IsLoopback() || ip.IsUnspecified() || !ip.IsGlobalUnicast() {
			continue
		}
		// Link-local IPv6 needs a zone to be dialable and rarely helps here.
		if ip.To4() == nil && ip.IsLinkLocalUnicast() {
			continue
		}
		endpoints = append(endpoints, net.JoinHostPort(ip.String(), fmt.Sprint(port)))
	}
	return endpoints
}

// dialDirect races the candidate addresses and returns the first connection
// that completes the handshake.
//
// Racing rather than trying in sequence matters: an unreachable address on a
// different subnet typically does not refuse the connection, it hangs until
// the timeout, and trying those one after another would take longer than the
// relay fallback it is meant to avoid.
func dialDirect(endpoints []string, tunnelID, peerToken string) (*directConn, string, error) {
	if len(endpoints) == 0 {
		return nil, "", fmt.Errorf("no candidate addresses")
	}

	greeting, err := buildHandshake(tunnelID, peerToken)
	if err != nil {
		return nil, "", err
	}

	type result struct {
		conn *directConn
		addr string
		err  error
	}
	results := make(chan result, len(endpoints))

	for _, endpoint := range endpoints {
		go func(addr string) {
			conn, err := attemptDirect(addr, greeting)
			results <- result{conn: conn, addr: addr, err: err}
		}(endpoint)
	}

	var lastErr error
	var winner *directConn
	var winnerAddr string

	// Collect every result so that a connection completing after we already
	// have a winner still gets closed instead of leaking.
	for range endpoints {
		r := <-results
		switch {
		case r.err != nil:
			lastErr = r.err
		case winner == nil:
			winner, winnerAddr = r.conn, r.addr
		default:
			r.conn.Close()
		}
	}

	if winner == nil {
		return nil, "", fmt.Errorf("no address reachable: %w", lastErr)
	}
	return winner, winnerAddr, nil
}

// attemptDirect performs one dial plus handshake.
func attemptDirect(addr string, greeting []byte) (*directConn, error) {
	conn, err := net.DialTimeout("tcp", addr, directDialTimeout)
	if err != nil {
		return nil, err
	}

	conn.SetDeadline(time.Now().Add(directHandshakeTimeout))

	if _, err := conn.Write(greeting); err != nil {
		conn.Close()
		return nil, fmt.Errorf("sending handshake to %s: %w", addr, err)
	}

	reply := make([]byte, handshakeReplySize)
	if _, err := io.ReadFull(conn, reply); err != nil {
		conn.Close()
		return nil, fmt.Errorf("reading handshake reply from %s: %w", addr, err)
	}
	if err := parseHandshakeReply(reply); err != nil {
		conn.Close()
		return nil, fmt.Errorf("handshake with %s: %w", addr, err)
	}

	// Clear the handshake deadline; tunnel traffic has no fixed timing.
	conn.SetDeadline(time.Time{})

	return newDirectConn(conn), nil
}