package proxy

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

	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/grpc/pb"
)

const (
	chunkSize         = 32 * 1024
	rendezvousTimeout = 10 * time.Second
)

// AgentStream is the server's half of a ProxyStream. It is an interface so the
// relay can be tested without gRPC.
type AgentStream interface {
	Send(*pb.ProxyServerMsg) error
	Recv() (*pb.ProxyClientMsg, error)
}

// Session owns one ephemeral listener and relays the single connection that
// arrives on it to an agent stream.
type Session struct {
	listener net.Listener
	allowed  []string
	timeout  time.Duration

	once    sync.Once
	mu      sync.Mutex
	reason  string
	conn    net.Conn
	closing bool
	onEnd   func(string)

	rendezvous *time.Timer
}

// OnEnd registers a callback invoked exactly once, when the session finishes
// tearing down, with the recorded reason (empty if it ended cleanly).
//
// It exists because the process that observes a relay failing is not
// necessarily the process that has to record it: with several replicas the
// listener lives on the pod holding the agent's stream, while the audit event
// belongs to the pod serving the browser's WebSocket.
func (s *Session) OnEnd(fn func(string)) {
	s.mu.Lock()
	s.onEnd = fn
	s.mu.Unlock()
}

// NewSession binds an ephemeral port on listenHost. allowed is the set of IPs
// permitted to connect; an empty set allows any, which is the degraded case
// when guacd's host could not be resolved.
//
// A watchdog is armed immediately: if the agent never opens its ProxyStream
// and calls Serve, nothing else would ever bound how long the listener (and
// the registry entry that references it) stays open. Serve cancels it once
// entered and re-arms the same deadline for the accept wait.
func NewSession(listenHost string, allowed []string) (*Session, error) {
	ln, err := net.Listen("tcp", net.JoinHostPort(listenHost, "0"))
	if err != nil {
		return nil, fmt.Errorf("bind relay listener: %w", err)
	}
	s := &Session{listener: ln, allowed: allowed, timeout: rendezvousTimeout}
	s.rendezvous = time.AfterFunc(rendezvousTimeout, func() {
		s.Close("agent_timeout")
	})
	return s, nil
}

func (s *Session) Port() int {
	return s.listener.Addr().(*net.TCPAddr).Port
}

// Reason reports why the session ended, empty if it ended cleanly.
func (s *Session) Reason() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.reason
}

// setReason records r as the session's failure reason, first write wins. It is
// a no-op once a deliberate teardown (Close) has begun: a local Close closing
// the conn out from under the relay goroutines produces exactly the kind of
// error (net.ErrClosed, a broken pipe on write, ...) that looks like a remote
// failure but is not one, and must not overwrite — or race to set — the real
// reason, or invent one where a clean local close has none.
func (s *Session) setReason(r string) {
	s.mu.Lock()
	if s.reason == "" && !s.closing {
		s.reason = r
	}
	s.mu.Unlock()
}

// Close tears the session down once. A non-empty reason is recorded only if no
// reason has been recorded already, and only before teardown begins. It closes
// both the listener and, if a connection has already been accepted, that
// connection too — an unconditional kill for the whole relay chain regardless
// of which stage it is in.
func (s *Session) Close(reason string) {
	if reason != "" {
		s.setReason(reason)
	}
	s.once.Do(func() {
		// Recorded before anything is actually closed: the reader/writer
		// goroutines in relay() call setReason from the errors this Close
		// itself is about to cause (a closed conn, a closed stream), and
		// those must be recognised as teardown noise, not a genuine failure.
		s.mu.Lock()
		s.closing = true
		conn := s.conn
		s.mu.Unlock()

		if s.rendezvous != nil {
			s.rendezvous.Stop()
		}
		_ = s.listener.Close()
		if conn != nil {
			_ = conn.Close()
		}

		// Read after teardown, not before: relay()'s goroutines set the reason
		// as they unwind, and reporting one recorded earlier than that would
		// name the symptom rather than the cause.
		s.mu.Lock()
		fn, final := s.onEnd, s.reason
		s.mu.Unlock()
		if fn != nil {
			fn(final)
		}
	})
}

func (s *Session) setConn(conn net.Conn) {
	s.mu.Lock()
	s.conn = conn
	s.mu.Unlock()
}

// Serve accepts exactly one connection, verifies its source, and relays until
// either side ends. It always closes the listener before returning.
func (s *Session) Serve(stream AgentStream) error {
	defer s.Close("")

	// The agent has claimed the session and opened its stream, so the
	// unclaimed-rendezvous watchdog no longer applies; the accept deadline set
	// just below takes over for the claimed case.
	if s.rendezvous != nil {
		s.rendezvous.Stop()
	}

	if l, ok := s.listener.(*net.TCPListener); ok {
		_ = l.SetDeadline(time.Now().Add(s.timeout))
	}

	// Accept until a connection from an allowed source arrives or the
	// rendezvous deadline (set once, above) expires. A rejected connection is
	// closed and the loop retries within the same deadline; the listener is
	// only closed once we have a valid connection, the deadline expires, or
	// Accept fails for a genuine (non-timeout) reason.
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			var netErr net.Error
			if errors.Is(err, os.ErrDeadlineExceeded) || (errors.As(err, &netErr) && netErr.Timeout()) {
				s.setReason("guacd_timeout")
			} else {
				s.setReason("accept_failed")
			}
			return fmt.Errorf("waiting for guacd: %w", err)
		}

		if !allowedRemote(conn.RemoteAddr().String(), s.allowed) {
			_ = conn.Close()
			continue
		}

		// One connection only: nothing else may claim this port. Store the
		// conn first so a concurrent external Close (e.g. the browser tab
		// closing) can reach it via the sync.Once teardown; then close just
		// the listener directly (not through Close, which would also close
		// the conn we are about to relay). The deferred s.Close("") above
		// performs the real one-shot teardown, including this conn, once
		// relay returns.
		s.setConn(conn)
		_ = s.listener.Close()

		return s.relay(conn, stream)
	}
}

func (s *Session) relay(conn net.Conn, stream AgentStream) error {
	errCh := make(chan error, 2)

	// guacd -> agent
	go func() {
		buf := make([]byte, chunkSize)
		for {
			n, err := conn.Read(buf)
			if n > 0 {
				chunk := make([]byte, n)
				copy(chunk, buf[:n])
				if sendErr := stream.Send(&pb.ProxyServerMsg{Data: chunk}); sendErr != nil {
					errCh <- sendErr
					return
				}
			}
			if err != nil {
				if !errors.Is(err, io.EOF) {
					s.setReason("guacd_read_error")
				}
				errCh <- err
				return
			}
		}
	}()

	// agent -> guacd
	go func() {
		for {
			msg, err := stream.Recv()
			if err != nil {
				errCh <- err
				return
			}
			if msg.Close != nil {
				s.setReason(msg.Close.Reason)
				errCh <- fmt.Errorf("agent closed relay: %s", msg.Close.Reason)
				return
			}
			if len(msg.Data) > 0 {
				if _, err := conn.Write(msg.Data); err != nil {
					errCh <- err
					return
				}
			}
		}
	}()

	err := <-errCh
	_ = conn.Close()
	if errors.Is(err, io.EOF) {
		return nil
	}
	return err
}

// allowedRemote reports whether remote (a host:port string) is in allowed. An
// empty allowed list permits anything.
func allowedRemote(remote string, allowed []string) bool {
	host, _, err := net.SplitHostPort(remote)
	if err != nil {
		log.Printf("proxy: unparseable remote address %q", remote)
		return false
	}
	if len(allowed) == 0 {
		return true
	}
	for _, a := range allowed {
		if a == host {
			return true
		}
	}
	return false
}
