package tunnel

// The relay half of the gateway (docs/relay-mode.md, issue #738 Phase 1):
// stable, keypair-claimed hostnames with a disconnect grace window. Trials
// (gateway.go) stay untouched; a relay differs in exactly three ways — its
// hostname belongs to its keypair, it has no 24-hour age-out, and a dropped
// connection puts the hostname to sleep instead of retiring it.

import (
	"context"
	"crypto/rand"
	"log"
	"math/big"
	"time"

	"github.com/coder/websocket"
)

func (g *Gateway) maxRelays() int { return orInt(g.MaxRelays, 200) }
func (g *Gateway) claimIdle() time.Duration {
	if g.ClaimIdle > 0 {
		return g.ClaimIdle
	}
	return 45 * 24 * time.Hour
}

// StartClaimJanitor frees claims whose owners haven't connected for ClaimIdle
// (default 45 days), hourly. Currently-connected tunnels are never swept, no
// matter how old their last_seen — Touch only runs at claim time.
func (g *Gateway) StartClaimJanitor() {
	if g.ClaimsDB == nil {
		return
	}
	go func() {
		for range time.Tick(time.Hour) {
			g.mu.Lock()
			keep := make(map[string]bool, len(g.tunnels))
			for h := range g.tunnels {
				keep[h] = true
			}
			g.mu.Unlock()
			freed, err := g.ClaimsDB.SweepIdle(g.claimIdle(), keep)
			if err != nil {
				log.Printf("tunnel: claim janitor: %v", err)
				continue
			}
			for _, h := range freed {
				log.Printf("tunnel: claim janitor freed %s (idle > %s)", h, g.claimIdle())
			}
		}
	}()
}
func (g *Gateway) relayGrace() time.Duration {
	if g.RelayGrace > 0 {
		return g.RelayGrace
	}
	return 15 * time.Minute
}

// handleRelayRegister runs after a hello with Mode "relay": challenge the
// client to prove its keypair, look up (or mint) its hostname, then serve the
// same probe/publish/hold flow trials get — minus the age-out, plus the
// sleep-on-disconnect grace.
func (g *Gateway) handleRelayRegister(ctx context.Context, c *websocket.Conn, ip string) {
	if g.RelayDomain == "" || g.ClaimsDB == nil {
		writeJSON(ctx, c, ControlMsg{Type: "error", Error: "this gateway doesn't offer relay hostnames"})
		c.Close(websocket.StatusPolicyViolation, "no relay support")
		return
	}

	nonce, err := randToken(24)
	if err != nil {
		return
	}
	if err := writeJSON(ctx, c, ControlMsg{Type: "challenge", Nonce: nonce}); err != nil {
		return
	}
	var claim ControlMsg
	if err := readJSON(ctx, c, &claim); err != nil || claim.Type != "claim" {
		writeJSON(ctx, c, ControlMsg{Type: "error", Error: "expected claim"})
		return
	}
	if err := VerifyClaimSig(claim.PubKey, nonce, claim.Sig); err != nil {
		writeJSON(ctx, c, ControlMsg{Type: "error", Error: "claim rejected: " + err.Error()})
		c.Close(websocket.StatusPolicyViolation, "bad claim")
		return
	}

	host, err := g.ClaimsDB.HostFor(claim.PubKey)
	if err != nil {
		writeJSON(ctx, c, ControlMsg{Type: "error", Error: "claims store unavailable"})
		return
	}
	if host == "" {
		// A brand-new claim pays the same admission toll as a trial; a
		// reconnect of an existing claim never does (a flaky home line must
		// not rate-limit itself out of its own hostname).
		if err := g.admit(ip); err != nil {
			writeJSON(ctx, c, ControlMsg{Type: "error", Error: err.Error()})
			return
		}
		count := 0
		g.ClaimsDB.sql.QueryRow(`SELECT COUNT(*) FROM relay_claims`).Scan(&count)
		if count >= g.maxRelays() {
			writeJSON(ctx, c, ControlMsg{Type: "error", Error: "this gateway has no free relay hostnames"})
			return
		}
		if host, err = g.newRelayHost(); err != nil {
			writeJSON(ctx, c, ControlMsg{Type: "error", Error: err.Error()})
			return
		}
		if err := g.ClaimsDB.Claim(host, claim.PubKey); err != nil {
			writeJSON(ctx, c, ControlMsg{Type: "error", Error: "couldn't record the claim, try again"})
			return
		}
		log.Printf("tunnel: relay %s claimed (from %s)", host, ip)
	}
	g.ClaimsDB.Touch(claim.PubKey)

	t := &liveTunnel{host: host, control: c, created: time.Now(), closed: make(chan struct{}),
		relay: true, pubkey: claim.PubKey}
	t.proxy = g.proxyFor(t)

	// Reattach: replace a sleeping (or superseded) tunnel for this host.
	g.mu.Lock()
	if old := g.tunnels[host]; old != nil {
		if old.grace != nil {
			old.grace.Stop()
		}
		// A previous live connection for the same key is superseded — the
		// newest device wins, matching how people actually move machines.
		go g.drop(old)
	}
	g.tunnels[host] = t
	g.mu.Unlock()

	if err := g.probeBackend(ctx, t); err != nil {
		msg := "only Eleven instances can be relayed: " + err.Error()
		writeJSON(ctx, c, ControlMsg{Type: "error", Error: msg})
		c.Close(websocket.StatusPolicyViolation, "refused")
		g.drop(t)
		log.Printf("tunnel: relay %s refused (probe): %v", host, err)
		return
	}
	if g.OnRegister != nil {
		if err := g.OnRegister(host); err != nil {
			writeJSON(ctx, c, ControlMsg{Type: "error", Error: "couldn't publish the address, try again"})
			g.drop(t)
			log.Printf("tunnel: relay register %s: %v", host, err)
			return
		}
	}
	log.Printf("tunnel: relay %s connected (from %s)", host, ip)
	if err := writeJSON(ctx, c, ControlMsg{Type: "ready", Host: host}); err != nil {
		g.sleepTunnel(t)
		return
	}

	// Hold: no age-out; re-probe like trials.
	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 {
		select {
		case <-readErr:
			g.sleepTunnel(t)
			log.Printf("tunnel: relay %s sleeping (disconnected; grace %s)", host, g.relayGrace())
			return
		case <-reprobe.C:
			if err := g.probeBackend(ctx, t); err != nil {
				probeFails++
				if probeFails >= 2 {
					writeJSON(ctx, c, ControlMsg{Type: "error", Error: "the relayed backend stopped looking like an Eleven instance"})
					g.drop(t)
					log.Printf("tunnel: relay %s dropped (re-probe): %v", host, err)
					return
				}
			} else {
				probeFails = 0
			}
		case <-t.closed:
			return
		case <-ctx.Done():
			g.sleepTunnel(t)
			return
		}
	}
}

// sleepTunnel keeps a relay's hostname routable-but-503 for the grace window,
// then retires it fully. The claim itself persists either way — the hostname
// still belongs to the keypair and comes back on the next connect.
func (g *Gateway) sleepTunnel(t *liveTunnel) {
	g.mu.Lock()
	if g.tunnels[t.host] != t { // superseded by a reattach
		g.mu.Unlock()
		return
	}
	t.sleeping = true
	t.grace = time.AfterFunc(g.relayGrace(), func() { g.drop(t) })
	g.mu.Unlock()
	t.control.CloseNow()
}

func (g *Gateway) newRelayHost() (string, error) {
	for range 20 {
		slug := make([]byte, 6)
		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 := "r-" + string(slug) + "." + g.RelayDomain
		if owner, err := g.ClaimsDB.Owner(host); err != nil {
			return "", err
		} else if owner == "" {
			return host, nil
		}
	}
	return "", errNoFreeHost
}