package tunnel

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"math/big"
	"net"
	"net/http"
	"net/http/httputil"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/coder/websocket"
)

// Gateway terminates the public side of every tunnel. It serves two kinds of
// host: its own control host (register/conn endpoints) and the try-* slug
// hosts, which it reverse-proxies down the matching tunnel. Route
// registration is delegated through OnRegister/OnClose so this package never
// touches the router's registry directly.
type Gateway struct {
	ControlHost string // e.g. "tunnel.11ch.at" — where clients register
	Domain      string // e.g. "11ch.at" — slugs become try-<slug>.<domain>

	// OnRegister makes the public hostname routable (a registry row pointing
	// back at this gateway's listener); OnClose removes it. Both must be safe
	// to call from multiple goroutines.
	OnRegister func(host string) error
	OnClose    func(host string)

	// Relay mode: stable keypair-claimed hostnames (docs/relay-mode.md).
	// Off unless RelayDomain and ClaimsDB are both set — the gateway then
	// refuses "relay" hellos, and trials behave exactly as before.
	RelayDomain string        // e.g. "example-relay.net" — relay hosts are r-<slug>.<RelayDomain>
	ClaimsDB    *Claims       // claim persistence (gateway-owned sqlite)
	MaxRelays   int           // claimed hostnames (0 = default 200)
	RelayGrace  time.Duration // route lifetime after a relay disconnects (0 = default 15m)
	ClaimIdle   time.Duration // free claims idle this long (0 = default 45 days)

	MaxTunnels  int           // concurrent tunnels (0 = default 50)
	MaxConns    int           // concurrent data connections per tunnel (0 = default 32)
	MaxPerIP    int           // new tunnels per IP per hour (0 = default 5)
	MaxAge      time.Duration // tunnel lifetime (0 = default 24h)
	DialTimeout time.Duration // client dial-back wait (0 = default 15s)

	mu      sync.Mutex
	tunnels map[string]*liveTunnel // by public host
	pending map[string]chan net.Conn
	byIP    map[string][]time.Time // register timestamps, pruned hourly
}

type liveTunnel struct {
	host    string
	control *websocket.Conn
	proxy   *httputil.ReverseProxy
	created time.Time
	closed  chan struct{}
	once    sync.Once
	// relay tunnels survive a disconnect: sleeping keeps the hostname
	// routable (answering 503) for RelayGrace so a home connection blip
	// isn't an outage; reattach cancels the grace timer.
	relay    bool
	pubkey   string
	sleeping bool
	grace    *time.Timer
	// conns counts live + in-flight data connections, bounding how much of the
	// trial machine's uplink one hostile visitor can claim.
	conns atomic.Int32
}

func (g *Gateway) maxTunnels() int { return orInt(g.MaxTunnels, 50) }
func (g *Gateway) maxConns() int   { return orInt(g.MaxConns, 32) }
func (g *Gateway) maxPerIP() int   { return orInt(g.MaxPerIP, 5) }
func (g *Gateway) maxAge() time.Duration {
	if g.MaxAge > 0 {
		return g.MaxAge
	}
	return 24 * time.Hour
}
func (g *Gateway) dialTimeout() time.Duration {
	if g.DialTimeout > 0 {
		return g.DialTimeout
	}
	return 15 * time.Second
}
func orInt(v, def int) int {
	if v > 0 {
		return v
	}
	return def
}

func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	host := r.Host
	if h, _, err := net.SplitHostPort(host); err == nil {
		host = h
	}
	if host == g.ControlHost {
		switch {
		case r.URL.Path == "/api/tunnel/register":
			g.handleRegister(w, r)
		case strings.HasPrefix(r.URL.Path, "/api/tunnel/conn/"):
			g.handleConn(w, r)
		case r.URL.Path == "/healthz":
			fmt.Fprintln(w, "ok")
		default:
			// A human wandered in; point them somewhere useful.
			http.Redirect(w, r, "https://elevenmessenger.com/self-host/", http.StatusFound)
		}
		return
	}
	g.mu.Lock()
	t := g.tunnels[host]
	g.mu.Unlock()
	if t == nil {
		http.Error(w, "this trial space isn't running — its lchat try command has stopped", http.StatusBadGateway)
		return
	}
	g.mu.Lock()
	asleep := t.sleeping
	g.mu.Unlock()
	if asleep {
		http.Error(w, "this space's machine is offline right now — it reconnects automatically when it's back", http.StatusServiceUnavailable)
		return
	}
	t.proxy.ServeHTTP(w, r)
}

// handleRegister accepts a client's control connection, assigns it a slug,
// makes the host routable, and then sits in a read loop until the client goes
// away (the read also services the client's keepalive pings).
func (g *Gateway) handleRegister(w http.ResponseWriter, r *http.Request) {
	ip := clientIP(r)
	if err := g.admit(ip); err != nil {
		http.Error(w, err.Error(), http.StatusTooManyRequests)
		return
	}
	c, err := websocket.Accept(w, r, nil)
	if err != nil {
		return
	}
	defer c.CloseNow()
	c.SetReadLimit(1 << 16)

	ctx := r.Context()
	var hello ControlMsg
	if err := readJSON(ctx, c, &hello); err != nil || hello.Type != "hello" {
		writeJSON(ctx, c, ControlMsg{Type: "error", Error: "expected hello"})
		return
	}

	if hello.Mode == "relay" {
		g.handleRelayRegister(ctx, c, ip)
		return
	}

	host, err := g.newHost()
	if err != nil {
		writeJSON(ctx, c, ControlMsg{Type: "error", Error: err.Error()})
		return
	}
	t := &liveTunnel{host: host, control: c, created: time.Now(), closed: make(chan struct{})}
	t.proxy = g.proxyFor(t)

	g.mu.Lock()
	g.tunnels[host] = t
	g.byIP[ip] = append(g.byIP[ip], time.Now())
	g.mu.Unlock()

	// The hostname goes public only after the backend proves it's an Eleven
	// instance (the client serves dial-backs from the moment it says hello).
	if err := g.probeBackend(ctx, t); err != nil {
		msg := "only Eleven instances can be tunnelled: " + err.Error()
		writeJSON(ctx, c, ControlMsg{Type: "error", Error: msg})
		// Graceful close so the error frame is delivered before the socket
		// dies (CloseNow would race it and the client would just see EOF).
		c.Close(websocket.StatusPolicyViolation, "refused")
		g.drop(t)
		log.Printf("tunnel: %s refused (probe): %v", host, err)
		return
	}

	if g.OnRegister != nil {
		if err := g.OnRegister(host); err != nil {
			g.drop(t)
			writeJSON(ctx, c, ControlMsg{Type: "error", Error: "couldn't publish the address, try again"})
			log.Printf("tunnel: register %s: %v", host, err)
			return
		}
	}
	log.Printf("tunnel: %s open (from %s)", host, ip)
	if err := writeJSON(ctx, c, ControlMsg{Type: "ready", Host: host}); err != nil {
		g.drop(t)
		return
	}

	// Hold the connection until the client leaves or the tunnel ages out. The
	// read loop keeps ping/pong serviced; clients never send data messages.
	ageOut := time.NewTimer(g.maxAge())
	defer ageOut.Stop()
	// Keep re-checking that the backend is still an Eleven instance; two
	// consecutive failures drop the tunnel (a client can't swap in a
	// different local server after passing the door check).
	reprobe := time.NewTicker(10 * time.Minute)
	defer reprobe.Stop()
	probeFails := 0
	readErr := make(chan error, 1)
	go func() {
		for {
			if _, _, err := c.Read(ctx); err != nil {
				readErr <- err
				return
			}
		}
	}()
	for done := false; !done; {
		select {
		case <-readErr:
			done = true
		case <-reprobe.C:
			if err := g.probeBackend(ctx, t); err != nil {
				probeFails++
				log.Printf("tunnel: %s re-probe failed (%d): %v", host, probeFails, err)
				if probeFails >= 2 {
					writeJSON(ctx, c, ControlMsg{Type: "error", Error: "the tunnelled backend stopped looking like an Eleven instance"})
					done = true
				}
			} else {
				probeFails = 0
			}
		case <-ageOut.C:
			writeJSON(ctx, c, ControlMsg{Type: "error", Error: "trial tunnel expired (24h) — run lchat try again for a fresh address"})
			done = true
		case <-t.closed:
			done = true
		case <-ctx.Done():
			done = true
		}
	}
	g.drop(t)
	log.Printf("tunnel: %s closed", host)
}

// handleConn is the client dialling back with the data connection for one
// pending id; it becomes a raw byte pipe paired with a waiting proxy dial.
func (g *Gateway) handleConn(w http.ResponseWriter, r *http.Request) {
	id := strings.TrimPrefix(r.URL.Path, "/api/tunnel/conn/")
	g.mu.Lock()
	ch := g.pending[id]
	delete(g.pending, id)
	g.mu.Unlock()
	if ch == nil {
		http.Error(w, "unknown conn id", http.StatusNotFound)
		return
	}
	c, err := websocket.Accept(w, r, nil)
	if err != nil {
		close(ch)
		return
	}
	c.SetReadLimit(-1)
	// Lifetime belongs to the proxied exchange now, not this handler's request
	// context — NetConn with Background keeps it open until either side closes.
	conn := websocket.NetConn(context.Background(), c, websocket.MessageBinary)
	select {
	case ch <- conn:
		// The proxy owns it now; block here isn't needed — Accept hijacked the
		// HTTP connection, and returning is fine.
	case <-time.After(5 * time.Second):
		conn.Close()
	}
}

// proxyFor builds the reverse proxy whose transport dials the client through
// the control channel: every new backend connection is one "dial" message and
// one dialled-back WebSocket.
func (g *Gateway) proxyFor(t *liveTunnel) *httputil.ReverseProxy {
	// A browser's first paint of the no-build shell is a burst of 40+ module
	// fetches over HTTP/2 — wider than the per-tunnel dial-back cap. Bounding
	// the transport just under the cap makes the overflow queue for a free
	// connection (as browsers always did against HTTP/1.1's 6-connection
	// limit) instead of erroring into 502s, any one of which permanently
	// fails the page's ES module graph. Two slots stay free for the gateway's
	// own liveness probes, which dial back outside this transport; the idle
	// pool matches so a finished request hands its connection to a queued one.
	perHost := max(1, g.maxConns()-2)
	transport := &http.Transport{
		DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
			return g.dialBack(ctx, t)
		},
		MaxConnsPerHost:     perHost,
		MaxIdleConns:        perHost,
		IdleConnTimeout:     60 * time.Second,
		MaxIdleConnsPerHost: perHost,
	}
	return &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			pr.SetXForwarded()
			pr.Out.URL.Scheme = "http"
			pr.Out.URL.Host = "tunnel-local"
			pr.Out.Host = pr.In.Host // the local instance sees its public hostname
		},
		Transport: transport,
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			http.Error(w, "the trial space didn't answer — is lchat try still running?", http.StatusBadGateway)
		},
	}
}

func (g *Gateway) dialBack(ctx context.Context, t *liveTunnel) (net.Conn, error) {
	if int(t.conns.Add(1)) > g.maxConns() {
		t.conns.Add(-1)
		return nil, fmt.Errorf("tunnel connection cap reached")
	}
	release := func() { t.conns.Add(-1) }
	id, err := randToken(18)
	if err != nil {
		release()
		return nil, err
	}
	ch := make(chan net.Conn, 1)
	g.mu.Lock()
	g.pending[id] = ch
	g.mu.Unlock()
	cleanup := func() {
		g.mu.Lock()
		delete(g.pending, id)
		g.mu.Unlock()
	}
	// The dial message rides the shared control channel, so its write is
	// bounded by the tunnel's own life, never by one visitor's request ctx:
	// the websocket library abandons the whole connection when a write is
	// cancelled mid-frame, which would let a single reload mid-load kill the
	// trial for everyone (writeJSON supplies its own 10s bound). A write
	// error therefore really does mean the control channel is dead.
	if err := writeJSON(context.Background(), t.control, ControlMsg{Type: "dial", ID: id}); err != nil {
		cleanup()
		release()
		g.drop(t) // control channel is dead — the tunnel is gone
		return nil, fmt.Errorf("tunnel control write: %w", err)
	}
	select {
	case conn, ok := <-ch:
		if !ok || conn == nil {
			cleanup()
			release()
			return nil, fmt.Errorf("tunnel dial-back failed")
		}
		return &countedConn{Conn: conn, release: release}, nil
	case <-time.After(g.dialTimeout()):
		cleanup()
		release()
		return nil, fmt.Errorf("tunnel dial-back timed out")
	case <-ctx.Done():
		cleanup()
		release()
		return nil, ctx.Err()
	}
}

// countedConn releases its tunnel's connection slot exactly once on close.
type countedConn struct {
	net.Conn
	release func()
	once    sync.Once
}

func (c *countedConn) Close() error {
	c.once.Do(c.release)
	return c.Conn.Close()
}

// probeBackend asks the tunnelled backend to prove it's an lchat instance
// before (and periodically after) its hostname goes public: GET /api/version
// through the tunnel must return 200 and a version. This keeps the trial
// domain lending its name — and its TLS — only to actual Eleven instances,
// not to whatever a modified client points --local at.
func (g *Gateway) probeBackend(ctx context.Context, t *liveTunnel) error {
	client := &http.Client{
		Timeout: 10 * time.Second,
		Transport: &http.Transport{
			DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
				return g.dialBack(ctx, t)
			},
		},
	}
	req, err := http.NewRequestWithContext(ctx, "GET", "http://tunnel-local/api/version", nil)
	if err != nil {
		return err
	}
	req.Host = t.host
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("backend didn't answer: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("backend /api/version: HTTP %d", resp.StatusCode)
	}
	var v struct {
		Version string `json:"version"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&v); err != nil || v.Version == "" {
		return fmt.Errorf("backend isn't an Eleven instance")
	}
	return nil
}

// drop removes the tunnel everywhere, once. The map guard matters for
// relays: a superseded tunnel dropping late must not take down the
// reattached one that now owns the hostname.
func (g *Gateway) drop(t *liveTunnel) {
	t.once.Do(func() {
		close(t.closed)
		g.mu.Lock()
		owns := g.tunnels[t.host] == t
		if owns {
			delete(g.tunnels, t.host)
		}
		g.mu.Unlock()
		if owns && g.OnClose != nil {
			g.OnClose(t.host)
		}
		t.control.CloseNow()
	})
}

var errNoFreeHost = fmt.Errorf("couldn't find a free address")

// admit enforces the per-IP rate and the global cap. Guardrails, not a wall:
// this is a free trial surface on a shared box, and the failure mode we're
// preventing is someone scripting a thousand tunnels, not a family trying it
// twice.
func (g *Gateway) admit(ip string) error {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.tunnels == nil {
		g.tunnels = map[string]*liveTunnel{}
		g.pending = map[string]chan net.Conn{}
		g.byIP = map[string][]time.Time{}
	}
	if len(g.tunnels) >= g.maxTunnels() {
		return fmt.Errorf("the trial gateway is full right now — please try again later")
	}
	cutoff := time.Now().Add(-time.Hour)
	kept := g.byIP[ip][:0]
	for _, ts := range g.byIP[ip] {
		if ts.After(cutoff) {
			kept = append(kept, ts)
		}
	}
	g.byIP[ip] = kept
	if len(kept) >= g.maxPerIP() {
		return fmt.Errorf("too many trial spaces from this address — please try again in an hour")
	}
	return nil
}

func (g *Gateway) newHost() (string, error) {
	for range 20 {
		slug := make([]byte, 5)
		for i := range slug {
			n, err := rand.Int(rand.Reader, big.NewInt(int64(len(slugAlphabet))))
			if err != nil {
				return "", err
			}
			slug[i] = slugAlphabet[n.Int64()]
		}
		host := "try-" + string(slug) + "." + g.Domain
		g.mu.Lock()
		_, taken := g.tunnels[host]
		g.mu.Unlock()
		if !taken {
			return host, nil
		}
	}
	return "", fmt.Errorf("couldn't find a free address")
}

// clientIP prefers the router's X-Forwarded-For (it fronts us over a unix
// socket, so RemoteAddr is meaningless).
func clientIP(r *http.Request) string {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		return strings.TrimSpace(strings.Split(xff, ",")[0])
	}
	if h, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		return h
	}
	return r.RemoteAddr
}

func randToken(n int) (string, error) {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

func readJSON(ctx context.Context, c *websocket.Conn, v any) error {
	_, data, err := c.Read(ctx)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, v)
}

func writeJSON(ctx context.Context, c *websocket.Conn, v any) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	wctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	return c.Write(wctx, websocket.MessageText, data)
}