package services

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"os"
	"sync"
	"time"

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

// ErrAgentOffline means the console cannot be opened because the target's agent
// is not on the command stream. Every console session is relayed by the agent,
// so this is fatal rather than a degraded mode.
var ErrAgentOffline = errors.New("agent is not connected")

// A console session spans up to three processes once there is more than one
// replica, and no two of them can be assumed to be the same one:
//
//	the browser's WebSocket   lands on an arbitrary pod
//	the agent's CommandStream lands on the pod holding presence for it
//	the agent's ProxyStream   lands on an arbitrary pod
//
// That third line is the one that is easy to get wrong. A ProxyStream is a
// separate HTTP/2 request, and an L7 proxy (Traefik, which the chart's gRPC
// ingress uses) balances requests rather than connections — so it does not
// follow the command stream. Binding the relay listener on the command
// stream's pod therefore fails roughly (n-1)/n of the time with "proxy session
// not found": the stream arrives at a pod whose registry is empty.
//
// So the listener is bound by whichever pod receives the ProxyStream, at the
// moment it receives it, and that pod announces its own address on
// ProxyAddrChannel. The WebSocket's pod subscribes before dispatching and
// hands the announced address to guacd. The address is the announcing pod's
// own, which is why it must resolve to a single pod (POD_IP under Kubernetes)
// rather than to the Service.
//
// Authorisation cannot live in that pod's memory either, so a pending record
// in Redis (bus.SetPendingProxy) carries the instance and server a proxy_id
// was minted for, and is consumed atomically on first claim.
//
// Teardown needs no message of its own. When the browser goes away guac closes
// its connection to the relay, the relay sees the read end, and the session
// closes itself — the same path a single-process deployment always took. Only
// the *reason* has to cross back, because the pod that writes the audit event
// is not the pod that observed the failure.

// How long Close waits for the relay's terminal reason to arrive before giving
// up. The audit event is written immediately afterwards; a reason that has not
// crossed the bus in this long is not going to improve the record by being
// waited for longer.
const proxyEndGrace = 2 * time.Second

// How long a minted proxy_id stays claimable, and how long the WebSocket's pod
// waits for the relay's address to be announced. The TTL is the longer of the
// two on purpose: a record that expired while its own opener was still waiting
// would turn a slow agent into an unexplained refusal.
const (
	proxyPendingTTL = 30 * time.Second
	proxyAddrWait   = 15 * time.Second
)

// ConsoleProxy is a relay as seen by the pod serving the WebSocket.
type ConsoleProxy struct {
	ProxyID string
	Host    string
	Port    int

	serverID string

	mu     sync.Mutex
	reason string
	closed bool
	stop   func()
	ended  chan struct{}
}

// proxyEnd is the terminal event a relay's owner pod publishes.
type proxyEnd struct {
	Reason string `json:"reason"`
}

func proxyListenHost() string {
	if v := os.Getenv("PROXY_LISTEN_HOST"); v != "" {
		return v
	}
	return "0.0.0.0"
}

// proxyAdvertiseHost is the address guacd will dial to reach a relay bound by
// *this* process. POD_IP wins over the configured value because with several
// replicas the configured value names the Service, and a Service cannot address
// one pod. The chart sets POD_IP from the downward API.
func proxyAdvertiseHost() string {
	if v := os.Getenv("POD_IP"); v != "" {
		return v
	}
	if v := os.Getenv("PROXY_ADVERTISE_HOST"); v != "" {
		return v
	}
	return "server"
}

func guacdAddr() string {
	if v := os.Getenv("GUACD_ADDR"); v != "" {
		return v
	}
	return "guacd:4822"
}

// guacdHosts resolves guacd's address to the IPs allowed to claim a relay
// listener. An unresolvable host yields an empty set, which allows any source:
// refusing everything would take the console down entirely, so the narrower
// protections (ephemeral port, 10s window, single accept) carry it instead.
func guacdHosts(addr string) []string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		host = addr
	}
	if ip := net.ParseIP(host); ip != nil {
		return []string{ip.String()}
	}
	ips, err := net.LookupHost(host)
	if err != nil {
		log.Printf("proxy: cannot resolve guacd host %q, allowing any relay source: %v", host, err)
		return nil
	}
	return ips
}

// proxyAddr is what a relay's binding pod announces: the address guacd should
// dial to reach the listener it has just bound.
type proxyAddr struct {
	Host string `json:"host"`
	Port int    `json:"port"`
}

// ClaimProxyStream authorises an incoming ProxyStream, binds a relay listener
// for it on this pod, and announces the address to whichever pod is serving the
// browser's WebSocket. It is called from the gRPC handler, on whichever replica
// the stream happened to reach.
//
// instanceID and serverID are the *authenticated* identity of the calling
// agent; they must match the pending record or the claim is refused, so an
// agent cannot relay a console session minted for another server.
func ClaimProxyStream(instanceID, serverID, proxyID string) (*proxy.Session, error) {
	ctx, cancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
	defer cancel()

	wantInstance, wantServer := bus.ClaimPendingProxy(ctx, proxyID)
	if wantInstance == "" {
		return nil, proxy.ErrNotFound
	}
	if wantInstance != instanceID || wantServer != serverID {
		return nil, proxy.ErrForbidden
	}

	sess, err := proxy.NewSession(proxyListenHost(), guacdHosts(guacdAddr()))
	if err != nil {
		return nil, err
	}

	sess.OnEnd(func(reason string) {
		endCtx, endCancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
		defer endCancel()
		if _, err := bus.Publish(endCtx, bus.ProxyEndChannel+proxyID, proxyEnd{Reason: reason}); err != nil {
			log.Printf("proxy: publish end for %s: %v", proxyID, err)
		}
	})

	addr := proxyAddr{Host: proxyAdvertiseHost(), Port: sess.Port()}

	// Logged on the success path, not just on failure. The address guacd is
	// about to be sent to is chosen per pod (POD_IP), so when a console fails
	// for some replicas and not others this line is the difference between
	// seeing which one answered and inferring it from silence.
	log.Printf("proxy %s (server %s): relay bound on %s:%d, node %s",
		proxyID, serverID, addr.Host, addr.Port, bus.NodeID())

	if _, err := bus.Publish(ctx, bus.ProxyAddrChannel+proxyID, addr); err != nil {
		// Nobody will ever dial this listener, so it is closed now rather than
		// left to sit out its rendezvous timeout.
		sess.Close("announce_failed")
		return nil, fmt.Errorf("announce relay address: %w", err)
	}

	return sess, nil
}

// OpenConsoleProxy asks the pod holding serverID's stream to bind a relay and
// tell the agent to meet it. The caller must Close the result.
func OpenConsoleProxy(instanceID, serverID string, targetPort int) (*ConsoleProxy, error) {
	if !Dispatcher.IsConnected(serverID) {
		return nil, ErrAgentOffline
	}

	proxyID, err := proxy.NewID()
	if err != nil {
		return nil, fmt.Errorf("generate proxy id: %w", err)
	}

	// Both subscriptions are established before the command is dispatched: a
	// fast agent binds and announces its relay within milliseconds, and a relay
	// that fails immediately publishes its reason just as quickly. Either
	// arriving before the subscriber is in place would be lost.
	cp := &ConsoleProxy{ProxyID: proxyID, serverID: serverID, ended: make(chan struct{})}
	endCtx, endCancel := context.WithCancel(context.Background())
	ends, unsubEnd, err := bus.Subscribe(endCtx, bus.ProxyEndChannel+proxyID)
	if err != nil {
		endCancel()
		return nil, fmt.Errorf("subscribe relay end: %w", err)
	}
	addrs, unsubAddr, err := bus.Subscribe(endCtx, bus.ProxyAddrChannel+proxyID)
	if err != nil {
		endCancel()
		unsubEnd()
		return nil, fmt.Errorf("subscribe relay address: %w", err)
	}
	cp.stop = func() {
		endCancel()
		unsubAddr()
		unsubEnd()
	}
	go cp.watchEnd(ends)

	// The pending record authorises the ProxyStream the agent is about to open,
	// and is written before the command so it cannot lose the race with it.
	pendCtx, pendCancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
	if err := bus.SetPendingProxy(pendCtx, proxyID, instanceID, serverID, proxyPendingTTL); err != nil {
		pendCancel()
		cp.stop()
		return nil, fmt.Errorf("register pending relay: %w", err)
	}
	pendCancel()

	abandon := func() {
		ctx, cancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
		bus.ClearPendingProxy(ctx, proxyID)
		cancel()
		cp.stop()
	}

	if _, err := Dispatcher.send(CommandEnvelope{
		ServerID: serverID,
		Command: &pb.ServerCommand{
			CommandId: proxyID,
			OpenProxy: &pb.OpenProxyCmd{ProxyId: proxyID, Port: uint32(targetPort)},
		},
	}); err != nil {
		abandon()
		if errors.Is(err, ErrAgentNotConnected) {
			return nil, ErrAgentOffline
		}
		return nil, err
	}

	// The command has reached the agent; the relay's address arrives only once
	// the agent has actually opened its ProxyStream somewhere in the fleet.
	var addr proxyAddr
	select {
	case b, ok := <-addrs:
		if !ok {
			abandon()
			return nil, fmt.Errorf("relay address subscription closed")
		}
		if err := json.Unmarshal(b, &addr); err != nil {
			abandon()
			return nil, fmt.Errorf("undecodable relay address: %w", err)
		}
	case <-time.After(proxyAddrWait):
		abandon()
		return nil, fmt.Errorf("agent did not open a relay for %s", serverID)
	}
	if addr.Host == "" || addr.Port == 0 {
		abandon()
		return nil, fmt.Errorf("relay opened without an address")
	}

	cp.Host = addr.Host
	cp.Port = addr.Port
	return cp, nil
}

func (c *ConsoleProxy) watchEnd(ends <-chan []byte) {
	defer close(c.ended)
	b, ok := <-ends
	if !ok {
		return
	}
	var e proxyEnd
	if err := json.Unmarshal(b, &e); err != nil {
		return
	}
	c.mu.Lock()
	if c.reason == "" {
		c.reason = e.Reason
	}
	c.mu.Unlock()
}

// Close waits briefly for the relay's terminal reason, then releases the
// subscription. It is safe to call more than once.
func (c *ConsoleProxy) Close() {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return
	}
	c.closed = true
	c.mu.Unlock()

	select {
	case <-c.ended:
	case <-time.After(proxyEndGrace):
	}
	c.stop()
}

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