package client

import (
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

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

// directListener accepts incoming direct tunnel connections.
//
// Only the share side listens: the use side is the one that knows a tunnel
// has been granted, so it makes the call. Tunnels are registered here as they
// are granted, and an incoming connection is matched against them.
type directListener struct {
	listener net.Listener
	secret   *crypto.TunnelSecret

	mu       sync.Mutex
	expected map[string]*expectedTunnel // tunnel ID -> pending acceptance
	closed   bool
}

// expectedTunnel is a granted tunnel waiting for its peer to connect.
type expectedTunnel struct {
	token    string
	accepted chan *directConn
}

// newDirectListener starts listening on the given port.
// Port 0 picks a free one, which is the sensible default: the actual port is
// advertised to peers, so it does not need to be predictable.
func newDirectListener(port int, secret *crypto.TunnelSecret) (*directListener, error) {
	if secret == nil {
		return nil, fmt.Errorf("direct connections require the tunnel secret")
	}

	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return nil, fmt.Errorf("listening for direct connections: %w", err)
	}

	dl := &directListener{
		listener: ln,
		secret:   secret,
		expected: make(map[string]*expectedTunnel),
	}

	go dl.acceptLoop()
	log.Printf("[direct] listening on %s", ln.Addr())

	return dl, nil
}

// Port returns the port actually bound.
func (dl *directListener) Port() int {
	if addr, ok := dl.listener.Addr().(*net.TCPAddr); ok {
		return addr.Port
	}
	return 0
}

// Expect registers a granted tunnel and returns a channel that receives the
// connection once a peer completes the handshake for it.
func (dl *directListener) Expect(tunnelID string) (<-chan *directConn, error) {
	accepted := make(chan *directConn, 1)

	dl.mu.Lock()
	defer dl.mu.Unlock()
	if dl.closed {
		return nil, fmt.Errorf("listener is closed")
	}
	dl.expected[tunnelID] = &expectedTunnel{
		token:    dl.secret.PeerToken(tunnelID),
		accepted: accepted,
	}
	return accepted, nil
}

// Forget drops a tunnel, whether it was taken over directly or fell back to
// the relay. Leaving entries behind would let a peer connect to a tunnel that
// is no longer live.
func (dl *directListener) Forget(tunnelID string) {
	dl.mu.Lock()
	defer dl.mu.Unlock()
	delete(dl.expected, tunnelID)
}

// Close stops accepting connections.
func (dl *directListener) Close() error {
	dl.mu.Lock()
	dl.closed = true
	dl.expected = make(map[string]*expectedTunnel)
	dl.mu.Unlock()
	return dl.listener.Close()
}

func (dl *directListener) acceptLoop() {
	for {
		conn, err := dl.listener.Accept()
		if err != nil {
			dl.mu.Lock()
			closed := dl.closed
			dl.mu.Unlock()
			if closed {
				return
			}
			// A transient accept error should not kill the listener, but it
			// must not spin either.
			log.Printf("[direct] accept error: %v", err)
			time.Sleep(100 * time.Millisecond)
			continue
		}

		go dl.handleIncoming(conn)
	}
}

// handleIncoming validates one incoming connection's handshake.
func (dl *directListener) handleIncoming(conn net.Conn) {
	conn.SetDeadline(time.Now().Add(directHandshakeTimeout))

	greeting := make([]byte, handshakeSize)
	if _, err := io.ReadFull(conn, greeting); err != nil {
		conn.Close()
		return
	}

	tunnelID, presented, err := parseHandshake(greeting)
	if err != nil {
		log.Printf("[direct] rejecting %s: %v", conn.RemoteAddr(), err)
		conn.Write(buildHandshakeReply(directWrongVerson))
		conn.Close()
		return
	}

	dl.mu.Lock()
	tunnel, known := dl.expected[tunnelID]
	dl.mu.Unlock()

	if !known {
		conn.Write(buildHandshakeReply(directUnknownTun))
		conn.Close()
		return
	}

	// The token proves group membership. The relay knows the tunnel ID — it
	// routed the grant — but cannot derive this, so it cannot impersonate a
	// peer, and neither can anything else that merely reaches the port.
	if !constantTimeEqual(presented, tunnel.token) {
		log.Printf("[direct] rejecting %s: bad token for tunnel", conn.RemoteAddr())
		conn.Write(buildHandshakeReply(directBadToken))
		conn.Close()
		return
	}

	if _, err := conn.Write(buildHandshakeReply(directAccepted)); err != nil {
		conn.Close()
		return
	}
	conn.SetDeadline(time.Time{})

	direct := newDirectConn(conn)

	// Claim the tunnel: whoever handshakes first wins, and a second connection
	// for the same tunnel is dropped rather than replacing a live one.
	dl.mu.Lock()
	current, still := dl.expected[tunnelID]
	if still && current == tunnel {
		delete(dl.expected, tunnelID)
	} else {
		still = false
	}
	dl.mu.Unlock()

	if !still {
		direct.Close()
		return
	}

	log.Printf("[direct] accepted connection from %s for tunnel %s", conn.RemoteAddr(), tunnelID)
	tunnel.accepted <- direct
}
