package main

import (
	"context"
	"crypto/tls"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"

	"golang.org/x/crypto/acme/autocert"

	"lchat/internal/activator"
	"lchat/internal/bot"
)

// router is the multi-instance front door: it terminates TLS (autocert) and
// reverse-proxies each configured host to its backend's unix socket. The
// host→backend map is swapped atomically, so reloads and zero-downtime deploys
// never drop a live connection. See docs/instance-router.md.
type router struct {
	reg       *Registry
	upstreams atomic.Pointer[map[string]upstream]
	// act, when non-nil, is the hibernation activator (activator.go):
	// hibernate-flagged instance routes stop being always-on systemd backends
	// and become on-demand child processes woken per request. Nil
	// (LCHAT_HIBERNATE_BUCKET unset) keeps today's proxy-only behavior; with it
	// set, unflagged routes are still proxied straight (the flag is per host).
	act *activatorSvc
	// fallbackURL, when set, is where requests for unknown hosts are redirected
	// (e.g. the apex marketing site) instead of a bare 404 — a wildcard DNS
	// record then makes typo'd or retired subdomains land somewhere useful. Note
	// the limit: over HTTPS this only helps hosts that still HAVE a cert (e.g. a
	// removed instance); a never-served host fails the TLS handshake first,
	// because HostPolicy refuses to mint certs for unknown hosts (anything else
	// would let strangers burn the ACME rate limit). The :80 path always works.
	fallbackURL string
}

// upstream is one served host: the proxy to its backend, plus — when the
// activator manages it — the route itself, so a request can trigger the wake.
type upstream struct {
	proxy *httputil.ReverseProxy
	route Route
	// hibernating marks an activator-managed instance: the router itself
	// knows whether its process is live (it's the router's child), so "dead
	// backend" detection is the activator's state table, not a failed dial.
	hibernating bool
}

func newRouter(reg *Registry, act *activatorSvc) (*router, error) {
	rt := &router{reg: reg, act: act}
	if err := rt.reload(); err != nil {
		return nil, err
	}
	return rt, nil
}

// reload rebuilds the host→backend map from the registry and swaps it in
// atomically. Safe to call while serving.
func (rt *router) reload() error {
	routes, err := rt.reg.list()
	if err != nil {
		return err
	}
	m := make(map[string]upstream, len(routes))
	// The hosts the activator should manage after this reload — feeds both the
	// per-upstream gate and the retire pass below.
	managed := make(map[string]bool)
	for _, t := range routes {
		// Only serve ready routes: a 'creating'/'failed' entry has no working
		// backend, so proxying it (and minting a cert for it) would just stall → 502.
		if t.Socket == "" || t.Status != "ready" {
			continue
		}
		// Hibernation is per-host: only an instance whose route carries the
		// hibernate flag is activator-managed; unflagged instances and sibling
		// services (home, push) are proxied straight. With the bucket env set
		// but no route flagged, the router therefore behaves exactly as it does
		// without the activator — so the binary + env can ship fleet-wide
		// first, and the flag flips one instance at a time.
		hib := rt.act != nil && t.Kind == "instance" && t.Hibernate
		if hib {
			managed[t.Host] = true
		}
		m[t.Host] = upstream{proxy: unixProxy(t.Socket), route: t, hibernating: hib}
	}
	rt.upstreams.Store(&m)
	// Reconcile the activator with the new route set: a host un-flagged (or
	// removed) since the last reload may still have a live child holding its
	// socket and S3 lease — retire it so a systemd unit can take over cleanly.
	// After the map swap above, so no new request re-enters the activator for a
	// host being retired.
	if rt.act != nil {
		rt.act.retireExcept(managed)
	}
	return nil
}

func (rt *router) current() map[string]upstream {
	if p := rt.upstreams.Load(); p != nil {
		return *p
	}
	return nil
}

// allowed reports whether host is configured — used as the autocert HostPolicy,
// so a certificate is only ever requested for a host we actually serve.
func (rt *router) allowed(host string) bool {
	_, ok := rt.current()[host]
	return ok
}

func (rt *router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	host := r.Host
	if h, _, err := net.SplitHostPort(host); err == nil {
		host = h
	}
	up, ok := rt.current()[host]
	if !ok {
		rt.unknownHost(w, r)
		return
	}
	if up.hibernating {
		// Wake (or join an in-progress wake) before proxying; the activator
		// also stamps this request's activity for the idle clock. Note a bare
		// TLS handshake never wakes anything — only a proxied request does.
		rt.act.serveHTTP(w, r, up.route, up.proxy)
		return
	}
	up.proxy.ServeHTTP(w, r)
}

// unknownHost answers a request for a host we don't serve: a redirect to the
// fallback (302 — a typo'd host must not be cached as permanently moved), or
// the classic 404 when no fallback is configured.
func (rt *router) unknownHost(w http.ResponseWriter, r *http.Request) {
	if rt.fallbackURL != "" {
		http.Redirect(w, r, rt.fallbackURL, http.StatusFound)
		return
	}
	http.Error(w, "unknown host", http.StatusNotFound)
}

// unixProxy reverse-proxies to a backend listening on a unix socket, preserving
// the client's Host header (so the backend sees the real hostname) and passing
// WebSocket upgrades straight through.
func unixProxy(socketPath string) *httputil.ReverseProxy {
	dialer := &net.Dialer{}
	return &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			pr.SetXForwarded()
			pr.Out.URL.Scheme = "http"
			pr.Out.URL.Host = "lchat" // placeholder; the unix dialer ignores it
			pr.Out.Host = pr.In.Host  // backend sees the real hostname
			// The sidecar marker is a router-edge signal only (activity
			// exclusion, activator.go); strip it so nothing downstream can
			// ever grow a meaning an outsider could spoof. The mesh's
			// forwarded marker (one-hop loop guard, mesh.go) gets the same
			// treatment for the same reason.
			pr.Out.Header.Del(bot.SidecarHeader)
			pr.Out.Header.Del(meshForwardedHeader)
		},
		Transport: &http.Transport{
			DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
				return dialUnixWithRetry(ctx, dialer, socketPath)
			},
			// Bound idle keep-alives so proxies replaced on reload don't leak their
			// pooled connections + reader goroutines on the long-running router.
			MaxIdleConns:    32,
			IdleConnTimeout: 90 * time.Second,
		},
	}
}

// dialUnixWithRetry masks the brief gap during an instance restart: while the old
// process has exited and the new one hasn't rebound the socket, dials fail with
// ENOENT/ECONNREFUSED. Retry those for a short window so no in-flight request is
// dropped; any other error fails fast. The window must exceed a deploy's
// socket-absent gap — the instance drain (10s in main.go) plus rebind — so an
// in-flight request is never 502'd mid-deploy; a sustained-down backend hangs this
// long before 502 (rare: systemd Restart=always makes "down" transient). True
// blue/green would remove the gap entirely (see docs/instance-router.md).
func dialUnixWithRetry(ctx context.Context, d *net.Dialer, socketPath string) (net.Conn, error) {
	const window = 13 * time.Second
	deadline := time.Now().Add(window)
	for {
		conn, err := d.DialContext(ctx, "unix", socketPath)
		if err == nil {
			return conn, nil
		}
		if ctx.Err() != nil || !isRetryableDialErr(err) || time.Now().After(deadline) {
			return nil, err
		}
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
}

func isRetryableDialErr(err error) bool {
	return errors.Is(err, syscall.ECONNREFUSED) || errors.Is(err, syscall.ENOENT)
}

// runRouter is the `lchat router` subcommand: the shared front door for every
// instance and sibling service on the box.
func runRouter(args []string) {
	fs := flag.NewFlagSet("router", flag.ExitOnError)
	registryPath := fs.String("registry", envOr("LCHAT_REGISTRY", "/data/registry.db"), "router registry SQLite path (env: LCHAT_REGISTRY)")
	acmeCache := fs.String("acme-cache", envOr("LCHAT_ACME_CACHE", "/data/acme"), "autocert certificate cache dir (env: LCHAT_ACME_CACHE)")
	httpsAddr := fs.String("https", envOr("LCHAT_HTTPS_ADDR", ":443"), "HTTPS listen address")
	httpAddr := fs.String("http", envOr("LCHAT_HTTP_ADDR", ":80"), "HTTP listen address (ACME HTTP-01 challenge + redirect to https)")
	fallbackURL := fs.String("fallback-url", os.Getenv("LCHAT_FALLBACK_URL"),
		"redirect requests for unknown hosts here, e.g. https://elevenmessenger.com (env: LCHAT_FALLBACK_URL); empty keeps the plain 404")
	wildcardDomains := fs.String("wildcard-domains", envOr("LCHAT_WILDCARD_DOMAINS", ""),
		"comma-separated base domains to serve a *.domain wildcard cert for, obtained via ACME DNS-01 (env: LCHAT_WILDCARD_DOMAINS; DNS-01 provider via LCHAT_DNS_PROVIDER — dnsimple (default, needs LCHAT_DNSIMPLE_TOKEN) or route53 (AWS credential chain))")
	_ = fs.Parse(args)

	reg, err := openRegistry(*registryPath)
	if err != nil {
		log.Fatalf("router: open registry %s: %v", *registryPath, err)
	}

	// Instance hibernation (activator.go): opt-in via LCHAT_HIBERNATE_BUCKET.
	// Unset keeps the classic shape — instances under systemd, router as a
	// pure proxy — so self-hosters and the current fleet are untouched.
	var act *activatorSvc
	hc, err := hibernateConfigFromEnv()
	if err != nil {
		log.Fatalf("router: %v", err)
	}
	// The activator mesh (mesh.go): opt-in via LCHAT_MESH_SELF, rides on the
	// hibernation bucket (meshFromEnv errors if it's set without one). Unset =
	// single box, no mesh codepaths.
	mesh, err := meshFromEnv(hc)
	if err != nil {
		log.Fatalf("router: %v", err)
	}
	if hc != nil {
		// materialize (cloud.go's controldMaterializeFunc) routes a route's
		// first-wake staging through the control daemon instead of writing
		// /opt/lchat directly, which this sandboxed process can't do (see
		// s3Backend.materialize's doc comment). nil — the common case, no
		// LoadCredential=control-token:… on this unit — falls back to the
		// direct calls, today's long-standing behavior, unchanged.
		materialize := controldMaterializeFunc()
		if materialize == nil {
			log.Printf("router: no control token available — a route materialized elsewhere (e.g. a mesh-replicated route) will fail its first wake here; see deploy/systemd/lchat-router-control-token.conf")
		}
		// The registry doubles as the due-time index's store (routes.next_due),
		// so scheduled deliveries survive router restarts and fire on time even
		// for sleeping instances.
		act = newActivator(*hc, &s3Backend{cfg: *hc, materialize: materialize}, reg)
		if mesh != nil {
			// Wired and beating before act.start(): the sweeper consults the
			// mesh from its first tick, and the first request deserves a
			// membership view.
			act.mesh = mesh
			// The registry is how the mesh reads a host's replicated
			// home_region for region-scoped placement (candidatesFor).
			mesh.reg = reg
			mesh.start()
			log.Printf("router: activator mesh on — this box is %s (region %q, heartbeats s3://%s/%s every %s, peers stale after %s, placement scoped to %v)",
				mesh.self, mesh.region, hc.Bucket, hc.meshPrefix(), mesh.interval, mesh.staleAfter, mesh.domains)

			// This box's own capacity snapshot for the Fleet status page's
			// multi-host aggregation (meshstats.go) — a separate S3 client on
			// the same store, since meshSvc's is interface-typed for tests.
			if statsStore, err := activator.NewHeartbeatStore(context.Background(), hc.Bucket, hc.meshPrefix(), hc.Region); err != nil {
				log.Printf("router: mesh stats reporter: %v — this box's capacity snapshot won't publish", err)
			} else {
				startStatsReporter(statsStore, mesh.self, reg, act, mesh)
			}
		}
		act.start()
		log.Printf("router: hibernation active (bucket %s, prefix %q, idle %s) — flagged instances wake on demand",
			hc.Bucket, hc.Prefix, hc.Idle)
	}

	rt, err := newRouter(reg, act)
	if err != nil {
		log.Fatalf("router: load routes: %v", err)
	}
	rt.fallbackURL = *fallbackURL

	// One cert cache (certcache.go) backs every issuance artifact — autocert's
	// per-host certs here AND the wildcard manager's PEMs/account key below —
	// so swapping the backend (LCHAT_CERT_CACHE) moves all of it at once.
	certCache, err := newCertCache(*acmeCache)
	if err != nil {
		log.Fatalf("router: %v", err)
	}

	m := &autocert.Manager{
		Prompt: autocert.AcceptTOS,
		Cache:  certCache,
		HostPolicy: func(_ context.Context, host string) error {
			if rt.allowed(host) {
				return nil
			}
			return fmt.Errorf("acme: host %q not in registry", host)
		},
	}

	// Wildcard certs for the space base domain(s) (wildcard.go): any host one
	// label under a configured domain is served the wildcard; everything else —
	// and those hosts too, until the first wildcard issuance lands — takes the
	// per-host autocert path above. Note the deliberate side effect: an unknown
	// *.11ch.at host now completes the handshake and gets the fallback/404
	// instead of a TLS error, since no per-host cert has to be minted first.
	tlsCfg := m.TLSConfig()
	if doms := splitDomains(*wildcardDomains); len(doms) > 0 {
		// DNSimple by default (LCHAT_DNSIMPLE_TOKEN), Route53 as the additive
		// off-ramp (LCHAT_DNS_PROVIDER=route53) — dnsProviderFromEnv decides. A
		// nil provider (missing token, bad selector) keeps per-host issuance.
		if dns, reason := dnsProviderFromEnv(); dns == nil {
			log.Printf("router: LCHAT_WILDCARD_DOMAINS set but no DNS-01 provider (%s) — wildcard certs stay off (per-host issuance continues)", reason)
		} else {
			wc := newWildcardManager(doms, certCache, dns)
			go wc.run()
			inner := tlsCfg.GetCertificate
			tlsCfg.GetCertificate = func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
				if cert := wc.certFor(hello.ServerName); cert != nil {
					return cert, nil
				}
				return inner(hello)
			}
		}
	}

	// Watch the registry for writes from other processes (the tunnel gateway
	// publishes routes this way, without needing rights to signal us). We poll
	// a cheap fingerprint of the routes table rather than PRAGMA data_version:
	// the pragma proved blind to cross-process commits on the production box
	// (a modernc/sqlite WAL-visibility quirk under the router's long-lived
	// connections), while ordinary SELECTs on this same handle demonstrably do
	// see fresh rows — reload()'s own queries are the proof. The fingerprint
	// moves on any insert, delete, or status/socket change; scanning a few
	// dozen rows every two seconds is noise.
	go func() {
		log.Printf("router: registry watcher armed (2s fingerprint poll)")
		// Each poll opens a FRESH connection: a brand-new connection reads the
		// latest committed state by construction, immune to whatever pins a
		// long-lived handle to a stale WAL snapshot (observed in production:
		// the shared handle went blind to cross-process commits — first with
		// PRAGMA data_version, then intermittently with plain SELECTs — while
		// separate processes saw fresh rows throughout). Opening SQLite costs
		// about a millisecond; at one poll every two seconds that is nothing.
		poll := func() (string, error) {
			db, err := sql.Open("sqlite", *registryPath)
			if err != nil {
				return "", err
			}
			defer db.Close()
			var fp string
			err = db.QueryRow(`SELECT COUNT(*) || ':' || COALESCE(MAX(rowid),0) || ':' ||
				COALESCE(SUM(rowid*17 + LENGTH(status)*3 + LENGTH(socket)),0) FROM routes`).Scan(&fp)
			return fp, err
		}
		// No baseline-only first tick: the first observed fingerprint also
		// triggers a reload, so a route published between the router's startup
		// reload and this poll (the tunnel gateway publishes exactly that way)
		// can't be silently baselined away. One redundant reload at startup is
		// the cost of closing that gap.
		var last string
		errs := 0
		for range time.Tick(2 * time.Second) {
			fp, err := poll()
			if err != nil {
				errs++
				if errs == 1 || errs%300 == 0 {
					log.Printf("router: registry-watch poll error (#%d): %v", errs, err)
				}
				continue
			}
			errs = 0
			if fp == last {
				continue
			}
			last = fp
			if err := rt.reload(); err != nil {
				log.Printf("router: registry-watch reload failed: %v", err)
				continue
			}
			log.Printf("router: registry changed — reloaded (%d route(s))", len(rt.current()))
		}
	}()

	// Reload routes on SIGHUP (add-instance/deploy will also trigger this via the
	// control socket in a later step).
	sighup := make(chan os.Signal, 1)
	signal.Notify(sighup, syscall.SIGHUP)
	go func() {
		for range sighup {
			if err := rt.reload(); err != nil {
				log.Printf("router: reload failed: %v", err)
				continue
			}
			log.Printf("router: reloaded (%d route(s))", len(rt.current()))
		}
	}()

	// The non-ACME :80 handler. Known hosts get autocert's classic redirect to
	// https (GET/HEAD only, 302 — mirrored here since supplying any fallback
	// replaces it); unknown hosts go to the fallback URL instead of a redirect
	// into a doomed TLS handshake (no cert will ever exist for them).
	httpFallback := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host := r.Host
		if h, _, err := net.SplitHostPort(host); err == nil {
			host = h
		}
		if !rt.allowed(host) {
			rt.unknownHost(w, r)
			return
		}
		if r.Method != "GET" && r.Method != "HEAD" {
			http.Error(w, "Use HTTPS", http.StatusBadRequest)
			return
		}
		http.Redirect(w, r, "https://"+host+r.URL.RequestURI(), http.StatusFound)
	})

	// The mesh route-replication endpoint (meshreplicate.go) only exists when
	// the mesh is on — everything else keeps proxying straight through rt, so
	// a non-mesh box's HTTPS handler is byte-identical to before this existed.
	var httpsHandler http.Handler = rt
	if mesh != nil {
		mux := http.NewServeMux()
		mux.Handle("POST /api/mesh/route", meshRouteHandler(reg, rt, mesh.token))
		// Deletion + reconciliation (#1346): a destroyed space's tombstone,
		// and the route listing mesh-seed --prune diffs against.
		mux.Handle("DELETE /api/mesh/route/{host}", meshRouteDeleteHandler(reg, rt, mesh.token))
		mux.Handle("GET /api/mesh/routes", meshRouteListHandler(reg, rt, mesh.token, mesh.domains))
		mux.Handle("/", rt)
		httpsHandler = mux
	}

	// ReadHeaderTimeout guards the public front door against Slowloris-style slow-
	// header connection exhaustion. No ReadTimeout/WriteTimeout — those would break
	// long-lived proxied WebSockets.
	httpSrv := &http.Server{Addr: *httpAddr, Handler: m.HTTPHandler(httpFallback), ReadHeaderTimeout: 10 * time.Second}
	httpsSrv := &http.Server{Addr: *httpsAddr, Handler: httpsHandler, TLSConfig: tlsCfg, ReadHeaderTimeout: 10 * time.Second}

	// :80 serves the ACME HTTP-01 challenge and redirects everything else to https.
	go func() {
		if err := httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("router: http listener %s: %v", *httpAddr, err)
		}
	}()
	go func() {
		log.Printf("router: serving %d route(s) on %s (registry %s, acme cache %s)",
			len(rt.current()), *httpsAddr, *registryPath, *acmeCache)
		if err := httpsSrv.ListenAndServeTLS("", ""); err != nil && err != http.ErrServerClosed {
			log.Fatalf("router: https listener %s: %v", *httpsAddr, err)
		}
	}()

	// Graceful shutdown on SIGTERM/SIGINT (e.g. a router binary upgrade): stop
	// accepting and drain in-flight HTTP before exiting. Proxied WebSockets are
	// hijacked (not drained); they reconnect through the new router.
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, syscall.SIGTERM, syscall.SIGINT)
	<-stop
	log.Printf("router: draining and shutting down")
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	_ = httpsSrv.Shutdown(ctx)
	_ = httpSrv.Shutdown(ctx)
	if act != nil {
		// Hibernated instances are the router's children: checkpoint each to
		// S3 and release its lease before exiting, so the next incarnation
		// wakes them cleanly instead of finding bound sockets and held leases.
		sctx, scancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer scancel()
		act.shutdown(sctx)
	}
}