// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

// Package control implements the control plane (spec §2): enrollment,
// profile, sessions. HTTPS with a possibly self-signed cert — clients trust
// the SPKI pin from enrollment, not a CA (spec §1).
package control

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"log/slog"
	"net"
	"net/http"
	"net/netip"
	"strconv"
	"strings"
	"time"

	"echo-lot.app/server/internal/session"
	"echo-lot.app/server/internal/store"
)

// Version is stamped by the build (-ldflags "-X ...").
var Version = "dev"

type Server struct {
	Store    *store.Store
	Sessions *session.Manager
	Name     string
	// Targets/capabilities for the profile response. The server offers only
	// what it actually implements; the registry grows with the code.
	UDPPort  int
	TCPPort  int
	StunPort int
	// SPKI pin of the serving cert, for the profile's pins[] field.
	PinB64 string
	// CertChain is the served leaf-first DER chain, for GET /v1/tls-reference.
	CertChain [][]byte
	// Capabilities as computed at startup from what is actually wired up.
	Capabilities []string
	// TCPRecent returns recent TCP-echo connections for a source IP (may be nil).
	TCPRecent func(ip string) any
	// DelayedEcho schedules/sends a DELAYED_ECHO for a session (may be nil).
	DelayedEcho func(sess *session.Session, actionID string) error
	// Granted server->client sends (spec §5). Both consume an asymmetric grant.
	DownTrain func(sess *session.Session, g *session.Grant, count, sizeBytes, intervalUs int) (int, error)
	BigSend   func(sess *session.Session, g *session.Grant, sizes []int) ([]int, error)
	// CanaryQueries returns logged canary lookups for a session prefix (may be nil).
	CanaryQueries func(sessionPrefix string) any
	// CanaryZone is surfaced in the profile so the app knows what to query.
	CanaryZone string
	// ProvenGood reports the server's self-test signal (may be nil). Surfaced
	// in the profile so a client can trust — or skip — MTU tests: if the
	// server's own egress isn't full-MTU, client MTU results measure the
	// server, not the client.
	ProvenGood func() (mtuOK, sysctlOK bool)
}

func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("POST /v1/enroll", s.enroll)
	mux.HandleFunc("GET /v1/profile", s.profile)
	mux.HandleFunc("POST /v1/sessions", s.newSession)
	mux.HandleFunc("DELETE /v1/sessions/{id}", s.deleteSession)
	mux.HandleFunc("GET /v1/sessions/{id}/observations", s.observations)
	mux.HandleFunc("POST /v1/sessions/{id}/actions", s.actions)
	mux.HandleFunc("POST /v1/echo", s.httpEcho)
	mux.HandleFunc("GET /v1/tls-reference", s.tlsReference)
	// TODO(spec §5): frag_send, throughput (both build on the same grant machinery)
	return mux
}

// EchoHandler exposes just the HTTP-echo endpoint for the optional cleartext
// listener (spec §4: plaintext-path tampering test).
func (s *Server) EchoHandler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("POST /v1/echo", s.httpEcho)
	return mux
}

// selftestSignal is the compact "server proven good" object for the profile.
// mtu_ok=false tells a client its MTU results would measure this server.
func selftestSignal(f func() (bool, bool)) map[string]any {
	if f == nil {
		return map[string]any{"mtu_ok": nil, "sysctl_ok": nil}
	}
	mtuOK, sysctlOK := f()
	return map[string]any{"mtu_ok": mtuOK, "sysctl_ok": sysctlOK}
}

// sessionAuth resolves {id} and requires the bearer to be the owning device.
func (s *Server) sessionAuth(w http.ResponseWriter, r *http.Request) *session.Session {
	dev := s.Store.DeviceByCredential(bearer(r))
	if dev == nil {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unknown credential"})
		return nil
	}
	sess := s.Sessions.ByID(r.PathValue("id"))
	if sess == nil || sess.Device != dev.ID {
		writeJSON(w, http.StatusNotFound, map[string]string{"error": "no such session"})
		return nil
	}
	return sess
}

// observations is spec §6 — everything the server witnessed for a session.
func (s *Server) observations(w http.ResponseWriter, r *http.Request) {
	sess := s.sessionAuth(w, r)
	if sess == nil {
		return
	}
	packetsSeen, udp, cb := sess.Observations()
	var tcp any
	if s.TCPRecent != nil {
		// Correlate by source IP: TCP echo carries no session id on the wire.
		if ds := sess.DataSource(); ds.IsValid() {
			tcp = s.TCPRecent(ds.Addr().Unmap().String())
		} else if sess.ControlSource.IsValid() {
			tcp = s.TCPRecent(sess.ControlSource.Unmap().String())
		}
	}
	var dnsCanary any
	if s.CanaryQueries != nil {
		dnsCanary = s.CanaryQueries(sess.ID[:16]) // the session's wire prefix
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"udp":          map[string]any{"packets_seen": packetsSeen, "packets": udp},
		"tcp":          tcp,
		"connect_back": cb,
		"dns_canary":   dnsCanary,
		// TODO(spec §6): http echo records
	})
}

// actions is spec §5 — authenticated asymmetric operations. Implemented:
// delayed_echo, connect_back. Destination is ALWAYS the session's observed
// source (data-plane source; connect_back uses the control-plane source).
func (s *Server) actions(w http.ResponseWriter, r *http.Request) {
	sess := s.sessionAuth(w, r)
	if sess == nil {
		return
	}
	var req struct {
		Action     string `json:"action"`
		DelayS     int    `json:"delay_s"`
		Protocol   string `json:"protocol"`
		Port       int    `json:"port"`
		Count      int    `json:"count"`
		SizeBytes  int    `json:"size_bytes"`
		IntervalUs int    `json:"interval_us"`
		SizesBytes []int  `json:"sizes_bytes"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "bad body"})
		return
	}
	actionID := randomID()
	switch req.Action {
	case "delayed_echo":
		if s.DelayedEcho == nil {
			writeJSON(w, http.StatusNotImplemented, map[string]string{"error": "delayed_echo not wired"})
			return
		}
		delay := min(max(req.DelayS, 1), 600)
		if !sess.DataSource().IsValid() {
			writeJSON(w, http.StatusConflict, map[string]string{"error": "no data-plane traffic seen yet — send an ECHO first"})
			return
		}
		time.AfterFunc(time.Duration(delay)*time.Second, func() {
			if err := s.DelayedEcho(sess, actionID); err != nil {
				slog.Debug("delayed echo failed", "err", err)
			}
		})
		writeJSON(w, http.StatusAccepted, map[string]any{"action_id": actionID, "delay_s": delay})
	case "connect_back":
		if req.Port < 1 || req.Port > 65535 || (req.Protocol != "tcp" && req.Protocol != "udp") {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "connect_back needs protocol tcp|udp and a port"})
			return
		}
		target := net.JoinHostPort(sess.ControlSource.Unmap().String(), strconv.Itoa(req.Port))
		go func() {
			start := time.Now()
			conn, err := net.DialTimeout(req.Protocol, target, 5*time.Second)
			res := session.ConnectBackResult{ActionID: actionID, RttMs: float64(time.Since(start).Microseconds()) / 1000}
			switch {
			case err == nil:
				res.Result = "connected"
				if req.Protocol == "udp" {
					// UDP "dial" always succeeds locally; send one datagram
					// so the client actually observes something.
					_, _ = conn.Write([]byte("echolot-connect-back " + actionID))
				}
				conn.Close()
			case isTimeout(err):
				res.Result = "timeout"
			default:
				res.Result = "refused"
			}
			sess.RecordConnectBack(res)
		}()
		writeJSON(w, http.StatusAccepted, map[string]any{"action_id": actionID, "target": target})
	case "downtrain":
		if s.DownTrain == nil {
			writeJSON(w, http.StatusNotImplemented, map[string]string{"error": "downtrain not wired"})
			return
		}
		// A downstream train sends far more than it receives, so it needs a grant (§3.4).
		count := clamp(req.Count, 1, 5000)
		size := clamp(req.SizeBytes, dataMinPacket, 1500)
		interval := clamp(req.IntervalUs, 0, 1_000_000)
		g := sess.NewGrant(actionID, int64(count*size), 0, session.DefaultGrantLimits)
		if g == nil {
			writeJSON(w, http.StatusConflict, noDataPlaneYet)
			return
		}
		go func() {
			sent, err := s.DownTrain(sess, g, count, size, interval)
			slog.Info("downtrain finished", "action", actionID, "sent", sent, "bytes", g.Sent(), "err", err)
		}()
		writeJSON(w, http.StatusAccepted, map[string]any{
			"action_id": actionID, "count": count, "size_bytes": size, "interval_us": interval,
			"grant": map[string]any{"max_bytes": g.MaxBytes, "max_kbps": g.MaxKbps},
		})

	case "big_send":
		if s.BigSend == nil {
			writeJSON(w, http.StatusNotImplemented, map[string]string{"error": "big_send not wired"})
			return
		}
		sizes := req.SizesBytes
		if len(sizes) == 0 {
			sizes = []int{1200, 1400, 1472, 1500, 2000, 4000}
		}
		if len(sizes) > 32 {
			sizes = sizes[:32]
		}
		total := 0
		for _, x := range sizes {
			total += clamp(x, dataMinPacket, 9000)
		}
		g := sess.NewGrant(actionID, int64(total), 0, session.DefaultGrantLimits)
		if g == nil {
			writeJSON(w, http.StatusConflict, noDataPlaneYet)
			return
		}
		go func() {
			attempted, err := s.BigSend(sess, g, sizes)
			slog.Info("big_send finished", "action", actionID, "attempted", attempted, "err", err)
		}()
		writeJSON(w, http.StatusAccepted, map[string]any{
			"action_id": actionID, "sizes_bytes": sizes,
			"grant": map[string]any{"max_bytes": g.MaxBytes, "max_kbps": g.MaxKbps},
		})

	default:
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "unknown or unimplemented action"})
	}
}

// dataMinPacket is the smallest datagram that still carries a header + a little payload.
const dataMinPacket = 40

var noDataPlaneYet = map[string]string{
	"error": "no data-plane traffic seen yet — send an ECHO first so the destination is verified",
}

func clamp(v, lo, hi int) int {
	if v < lo {
		return lo
	}
	if v > hi {
		return hi
	}
	return v
}

func isTimeout(err error) bool {
	var ne net.Error
	return errors.As(err, &ne) && ne.Timeout()
}

func randomID() string {
	var b [8]byte
	_, _ = rand.Read(b[:])
	return hex.EncodeToString(b[:])
}

func bearer(r *http.Request) string {
	h := r.Header.Get("Authorization")
	if v, ok := strings.CutPrefix(h, "Bearer "); ok {
		return strings.TrimSpace(v)
	}
	return ""
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(v)
}

// enroll redeems a single-use enrollment token for a device credential (§2.1).
func (s *Server) enroll(w http.ResponseWriter, r *http.Request) {
	tok := bearer(r)
	if tok == "" {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "missing enrollment token"})
		return
	}
	var body struct {
		Name string `json:"name"`
	}
	_ = json.NewDecoder(r.Body).Decode(&body) // body is optional
	dev, err := s.Store.Redeem(tok, body.Name)
	if err != nil {
		writeJSON(w, http.StatusForbidden, map[string]string{"error": err.Error()})
		return
	}
	slog.Info("device enrolled", "device", dev.ID, "name", dev.Name)
	writeJSON(w, http.StatusCreated, map[string]string{
		"device_id":  dev.ID,
		"credential": dev.Credential, // returned exactly once
	})
}

// profile is spec §2.2. Only implemented capabilities are advertised.
func (s *Server) profile(w http.ResponseWriter, r *http.Request) {
	dev := s.Store.DeviceByCredential(bearer(r))
	if dev == nil {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unknown credential"})
		return
	}
	host := r.Host
	if i := strings.LastIndex(host, ":"); i > 0 {
		host = host[:i]
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"profile_version": 1,
		"name":            s.Name,
		"server_version":  Version,
		// source_url makes GPL §6 compliance mechanical for operators of
		// modified builds and gives clients provenance for the measurement.
		"source_url":   "", // TODO: stamp from build metadata
		"capabilities": s.Capabilities,
		"targets": []map[string]any{{
			"id":        s.Name,
			"ip4":       host, // TODO: explicit configured addresses, v6, second STUN addr
			"udp_port":  s.UDPPort,
			"tcp_port":  s.TCPPort,
			"stun_port": s.StunPort,
		}},
		"pins":            []string{"pin-sha256:" + s.PinB64},
		"next_pins":       []string{},
		"canary_zone":     s.CanaryZone,
		"server_selftest": selftestSignal(s.ProvenGood),
		"limits":          map[string]any{"max_kbps": 50000, "max_session_s": 900},
	})
}

// newSession is spec §2.4.
func (s *Server) newSession(w http.ResponseWriter, r *http.Request) {
	cred := bearer(r)
	dev := s.Store.DeviceByCredential(cred)
	if dev == nil {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unknown credential"})
		return
	}
	srcHost, _, _ := strings.Cut(r.RemoteAddr, ":")
	src, _ := netip.ParseAddr(strings.Trim(srcHost, "[]"))
	sess, salt, err := s.Sessions.New(dev.ID, cred, src)
	if err != nil {
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "session creation failed"})
		return
	}
	writeJSON(w, http.StatusCreated, map[string]any{
		"session_id": sess.ID,
		"key_salt":   base64.StdEncoding.EncodeToString(salt),
		"epoch":      sess.Epoch.Format(time.RFC3339),
		"expires_s":  int(time.Until(sess.Expires).Seconds()),
	})
}

func (s *Server) deleteSession(w http.ResponseWriter, r *http.Request) {
	if s.Store.DeviceByCredential(bearer(r)) == nil {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unknown credential"})
		return
	}
	if id := r.PathValue("id"); len(id) >= 16 {
		s.Sessions.Delete(id)
	}
	w.WriteHeader(http.StatusNoContent)
}

// SpkiPinB64 computes the RFC 7469 pin (SHA-256 over the SPKI, base64) of the
// leaf certificate — what enrollment QR codes carry and profiles re-state.
func SpkiPinB64(cert tls.Certificate) (string, error) {
	leaf := cert.Leaf
	if leaf == nil {
		parsed, err := x509.ParseCertificate(cert.Certificate[0])
		if err != nil {
			return "", err
		}
		leaf = parsed
	}
	sum := sha256.Sum256(leaf.RawSubjectPublicKeyInfo)
	return base64.StdEncoding.EncodeToString(sum[:]), nil
}