package main

import (
	"bytes"
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"sync/atomic"
	"time"

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

// Wildcard TLS (ACME DNS-01) for the space base domain(s).
//
// Every new space at <sqid>.11ch.at used to mint its own per-host certificate
// (autocert HTTP-01), which has a hard Let's Encrypt ceiling (~50 new certs a
// week per registered domain) and means a never-served subdomain fails the TLS
// handshake outright. A single wildcard cert for *.<domain>, obtained over
// DNS-01 (a TXT record at _acme-challenge.<domain> via the DNSimple API),
// covers every space instantly and forever-many of them.
//
// The wildcard only changes how a handshake GETS a certificate: routing and
// the autocert HostPolicy are untouched, and per-host HTTP-01 remains the path
// for every host outside the wildcard domains (and the fallback for hosts
// inside them until the wildcard is first obtained — so a DNS-provider or ACME
// outage degrades to exactly today's behaviour, never a dead router).
//
// The DNS-01 solver is PLUGGABLE (the dnsProvider seam): DNSimple stays the
// default, fully-supported backend (LCHAT_DNSIMPLE_TOKEN), and Route53
// (route53dns.go) is an ADDITIVE off-ramp selected by LCHAT_DNS_PROVIDER=route53
// for boxes that would rather authenticate with an AWS IAM role than a token on
// disk (docs/multi-edge.md). Selection lives in dnsProviderFromEnv; everything
// below is provider-blind.
//
// Off unless LCHAT_WILDCARD_DOMAINS is set and a DNS provider is configured
// (DNSimple by default, so LCHAT_DNSIMPLE_TOKEN). See
// docs/instance-router.md § Wildcard TLS.

const letsEncryptDirectory = "https://acme-v02.api.letsencrypt.org/directory"

// wildcardRenewBefore is how close to expiry we re-issue. LE certs live 90
// days; renewing with 30 left tolerates a month of provider outage.
const wildcardRenewBefore = 30 * 24 * time.Hour

// dnsProvider is the DNS-01 seam. It publishes the ACME challenge TXT and hands
// back a closure that removes it — a closure so each backend keeps its own
// teardown detail private (DNSimple deletes by the numeric record id it minted;
// Route53 re-submits the record set). wildcardManager.solveDNS01 only ever calls
// present, so it never learns which backend is in use.
//
// Implementations: dnsimpleDNS (the default, this file) and route53DNS
// (route53dns.go, the additive off-ramp). Choose with dnsProviderFromEnv.
type dnsProvider interface {
	// present creates a TXT record named name under zone (e.g.
	// "_acme-challenge" under "11ch.at") holding value, returning a cleanup that
	// deletes it. cleanup runs on its own context (see solveDNS01's defer) so it
	// still fires after the issuance context is dead.
	present(ctx context.Context, zone, name, value string) (cleanup func(context.Context) error, err error)
}

// dnsProviderFromEnv selects the DNS-01 backend, keeping DNSimple the DEFAULT
// with its trigger unchanged:
//
//   - unset / LCHAT_DNS_PROVIDER=dnsimple → DNSimple, exactly as before — needs
//     LCHAT_DNSIMPLE_TOKEN (an account API token); a missing token leaves the
//     wildcard off (per-host issuance continues), the historical behaviour.
//   - LCHAT_DNS_PROVIDER=route53 → Route53 (the multi-edge off-ramp), authed by
//     the standard AWS credential chain (env vars / EC2-or-container IAM role) —
//     no DNS token on disk. The dev boxes use their zone-scoped role for
//     *.dev.11ch.at (docs/multi-edge.md).
//
// Returns (nil, reason) when nothing is configured, so the caller logs the human
// reason and falls back to per-host issuance rather than failing to boot.
func dnsProviderFromEnv() (dnsProvider, string) {
	switch strings.ToLower(strings.TrimSpace(os.Getenv("LCHAT_DNS_PROVIDER"))) {
	case "route53":
		r53, err := newRoute53DNS(context.Background())
		if err != nil {
			return nil, fmt.Sprintf("LCHAT_DNS_PROVIDER=route53 but AWS config load failed: %v", err)
		}
		return r53, ""
	case "", "dnsimple":
		token := os.Getenv("LCHAT_DNSIMPLE_TOKEN")
		if token == "" {
			return nil, "LCHAT_DNSIMPLE_TOKEN missing"
		}
		return newDNSimpleDNS(token), ""
	default:
		return nil, fmt.Sprintf("unknown LCHAT_DNS_PROVIDER %q (want dnsimple or route53)", os.Getenv("LCHAT_DNS_PROVIDER"))
	}
}

type wildcardManager struct {
	domains []string
	// cache is the shared cert cache (certcache.go) — the same backend
	// autocert uses for per-host certs, so pooled routers share the wildcard
	// too. Under the dir backend the keys land as the historical file names
	// beside autocert's certs (wildcard-<domain>.pem, wildcard-account.key).
	cache     autocert.Cache
	dns       dnsProvider
	directory string
	// One atomically-swapped cert per domain, read on every 11ch.at handshake.
	certs map[string]*atomic.Pointer[tls.Certificate]
}

func newWildcardManager(domains []string, cache autocert.Cache, dns dnsProvider) *wildcardManager {
	w := &wildcardManager{
		domains:   domains,
		cache:     cache,
		dns:       dns,
		directory: envOr("LCHAT_ACME_DIRECTORY", letsEncryptDirectory),
		certs:     map[string]*atomic.Pointer[tls.Certificate]{},
	}
	// The startup load is synchronous (runRouter wires certFor into the TLS
	// config right after this), so bound it: a dir read is instant, but a
	// remote cache backend must degrade to "re-obtain" rather than hang boot.
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	for _, d := range domains {
		w.certs[d] = &atomic.Pointer[tls.Certificate]{}
		raw, err := w.cache.Get(ctx, wildcardPEMKey(d))
		if errors.Is(err, autocert.ErrCacheMiss) {
			continue // first run for this domain — run() obtains one
		}
		if err == nil {
			var cert *tls.Certificate
			if cert, err = parseWildcardPEM(raw); err == nil {
				w.certs[d].Store(cert)
				log.Printf("router: wildcard *.%s loaded from cache (valid until %s)", d, cert.Leaf.NotAfter.Format(time.RFC3339))
				continue
			}
		}
		log.Printf("router: wildcard *.%s: cached cert unreadable (%v) — will re-obtain", d, err)
	}
	return w
}

// wildcardPEMKey is the cache key for a domain's wildcard cert+key PEM. It is
// deliberately the pre-seam file name, so the dir backend reads the PEMs an
// older binary wrote to the acme cache dir with no migration.
func wildcardPEMKey(domain string) string {
	return "wildcard-" + domain + ".pem"
}

// certFor returns the wildcard cert for serverName, or nil when the name isn't
// exactly one label under a configured domain (a wildcard doesn't cover the
// apex or deeper levels) or no usable cert is loaded — nil sends the handshake
// to the per-host autocert path.
func (w *wildcardManager) certFor(serverName string) *tls.Certificate {
	sn := strings.ToLower(strings.TrimSuffix(serverName, "."))
	for _, d := range w.domains {
		sub, ok := strings.CutSuffix(sn, "."+d)
		if !ok || sub == "" || strings.Contains(sub, ".") {
			continue
		}
		cert := w.certs[d].Load()
		if cert == nil || time.Now().After(cert.Leaf.NotAfter) {
			// Never serve a dead cert: a wildcard whose renewal has been broken
			// past its actual expiry falls back to per-host issuance (today's
			// behaviour) instead of hard-failing every space handshake.
			return nil
		}
		return cert
	}
	return nil
}

// run keeps the wildcard certs current: an immediate check at startup (first
// issuance), then every 12h, dropping to hourly retries after a failure.
func (w *wildcardManager) run() {
	for {
		delay := 12 * time.Hour
		for _, d := range w.domains {
			if !w.needsObtain(d, time.Now()) {
				continue
			}
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
			err := w.obtain(ctx, d)
			cancel()
			if err != nil {
				log.Printf("router: wildcard *.%s: %v — per-host issuance keeps covering; retrying in 1h", d, err)
				delay = time.Hour
				continue
			}
			cert := w.certs[d].Load()
			log.Printf("router: wildcard *.%s obtained (valid until %s)", d, cert.Leaf.NotAfter.Format(time.RFC3339))
		}
		time.Sleep(delay)
	}
}

func (w *wildcardManager) needsObtain(domain string, now time.Time) bool {
	cert := w.certs[domain].Load()
	return cert == nil || now.Add(wildcardRenewBefore).After(cert.Leaf.NotAfter)
}

// obtain runs one full ACME DNS-01 issuance for *.domain and swaps the result in.
func (w *wildcardManager) obtain(ctx context.Context, domain string) error {
	acctKey, err := w.accountKey(ctx)
	if err != nil {
		return fmt.Errorf("account key: %v", err)
	}
	client := &acme.Client{Key: acctKey, DirectoryURL: w.directory}
	if _, err := client.Register(ctx, &acme.Account{}, acme.AcceptTOS); err != nil && err != acme.ErrAccountAlreadyExists {
		return fmt.Errorf("acme register: %v", err)
	}

	order, err := client.AuthorizeOrder(ctx, []acme.AuthzID{{Type: "dns", Value: "*." + domain}})
	if err != nil {
		return fmt.Errorf("authorize order: %v", err)
	}
	for _, zurl := range order.AuthzURLs {
		if err := w.solveDNS01(ctx, client, domain, zurl); err != nil {
			return err
		}
	}

	certKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return err
	}
	csr, err := x509.CreateCertificateRequest(rand.Reader, &x509.CertificateRequest{
		Subject:  pkix.Name{CommonName: "*." + domain},
		DNSNames: []string{"*." + domain},
	}, certKey)
	if err != nil {
		return fmt.Errorf("csr: %v", err)
	}
	ders, _, err := client.CreateOrderCert(ctx, order.FinalizeURL, csr, true)
	if err != nil {
		return fmt.Errorf("finalize: %v", err)
	}
	leaf, err := x509.ParseCertificate(ders[0])
	if err != nil {
		return fmt.Errorf("parse issued cert: %v", err)
	}
	cert := &tls.Certificate{Certificate: ders, PrivateKey: certKey, Leaf: leaf}
	raw, err := encodeWildcardPEM(cert, certKey)
	if err != nil {
		return fmt.Errorf("encode: %v", err)
	}
	if err := w.cache.Put(ctx, wildcardPEMKey(domain), raw); err != nil {
		return fmt.Errorf("persist: %v", err)
	}
	w.certs[domain].Store(cert)
	return nil
}

// solveDNS01 answers one authorization: publish the challenge TXT record,
// wait for the CA to see it, and always clean the record up afterwards.
func (w *wildcardManager) solveDNS01(ctx context.Context, client *acme.Client, domain, authzURL string) error {
	authz, err := client.GetAuthorization(ctx, authzURL)
	if err != nil {
		return fmt.Errorf("get authorization: %v", err)
	}
	if authz.Status == acme.StatusValid {
		return nil // a recent order already proved this name
	}
	var chal *acme.Challenge
	for _, c := range authz.Challenges {
		if c.Type == "dns-01" {
			chal = c
			break
		}
	}
	if chal == nil {
		return errors.New("CA offered no dns-01 challenge")
	}
	txt, err := client.DNS01ChallengeRecord(chal.Token)
	if err != nil {
		return err
	}

	cleanup, err := w.dns.present(ctx, domain, "_acme-challenge", txt)
	if err != nil {
		return fmt.Errorf("publish challenge TXT: %v", err)
	}
	// Clean up with a fresh context: the record must not outlive the attempt
	// even when ctx is already dead (stale challenge TXTs break the NEXT run).
	defer func() {
		cctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		if err := cleanup(cctx); err != nil {
			log.Printf("router: wildcard *.%s: cleanup of challenge TXT failed: %v", domain, err)
		}
	}()

	// Best-effort: give the record a moment to reach the authoritative NS
	// before telling the CA to look. Proceed on timeout — the CA does its own
	// retries, and a hard failure just comes back through the retry loop.
	waitForTXT(ctx, "_acme-challenge."+domain, txt, time.Minute)

	if _, err := client.Accept(ctx, chal); err != nil {
		return fmt.Errorf("accept challenge: %v", err)
	}
	if _, err := client.WaitAuthorization(ctx, authz.URI); err != nil {
		return fmt.Errorf("authorization: %v", err)
	}
	return nil
}

// wildcardAccountKeyName is the ACME account key's cache key — the pre-seam
// file name, for the same no-migration reason as wildcardPEMKey.
const wildcardAccountKeyName = "wildcard-account.key"

// accountKey loads (or mints) the ACME account key. Separate from autocert's
// account on purpose: this flow manages its own registration lifecycle.
func (w *wildcardManager) accountKey(ctx context.Context) (*ecdsa.PrivateKey, error) {
	raw, err := w.cache.Get(ctx, wildcardAccountKeyName)
	if err == nil {
		block, _ := pem.Decode(raw)
		if block == nil {
			return nil, fmt.Errorf("%s: not PEM", wildcardAccountKeyName)
		}
		return x509.ParseECPrivateKey(block.Bytes)
	}
	if !errors.Is(err, autocert.ErrCacheMiss) {
		// Only a genuine miss mints a key: treating a backend fault as "no key
		// yet" would register a fresh ACME account on every hiccup. Fail and
		// let obtain()'s hourly retry come back around.
		return nil, err
	}
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	der, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, err
	}
	// Crash-safety (temp + rename, 0600) is the cache backend's contract now —
	// the dir backend inherits it from autocert.DirCache.
	raw = pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: der})
	if err := w.cache.Put(ctx, wildcardAccountKeyName, raw); err != nil {
		return nil, err
	}
	return key, nil
}

// waitForTXT polls DNS until name resolves to a TXT containing value, or the
// deadline passes (best-effort — see caller).
func waitForTXT(ctx context.Context, name, value string, timeout time.Duration) {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) && ctx.Err() == nil {
		if txts, err := net.DefaultResolver.LookupTXT(ctx, name); err == nil {
			for _, t := range txts {
				if t == value {
					return
				}
			}
		}
		time.Sleep(3 * time.Second)
	}
}

// ---- PEM encode/parse (key + chain in one blob; the cert cache owns
// placement, permissions, and write atomicity) ----------------------------------

func encodeWildcardPEM(cert *tls.Certificate, key *ecdsa.PrivateKey) ([]byte, error) {
	var buf bytes.Buffer
	der, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, err
	}
	if err := pem.Encode(&buf, &pem.Block{Type: "EC PRIVATE KEY", Bytes: der}); err != nil {
		return nil, err
	}
	for _, c := range cert.Certificate {
		if err := pem.Encode(&buf, &pem.Block{Type: "CERTIFICATE", Bytes: c}); err != nil {
			return nil, err
		}
	}
	return buf.Bytes(), nil
}

func parseWildcardPEM(raw []byte) (*tls.Certificate, error) {
	cert := &tls.Certificate{}
	for {
		var block *pem.Block
		block, raw = pem.Decode(raw)
		if block == nil {
			break
		}
		switch block.Type {
		case "EC PRIVATE KEY":
			key, err := x509.ParseECPrivateKey(block.Bytes)
			if err != nil {
				return nil, err
			}
			cert.PrivateKey = key
		case "CERTIFICATE":
			cert.Certificate = append(cert.Certificate, block.Bytes)
		}
	}
	if cert.PrivateKey == nil || len(cert.Certificate) == 0 {
		return nil, errors.New("missing key or certificate in PEM")
	}
	leaf, err := x509.ParseCertificate(cert.Certificate[0])
	if err != nil {
		return nil, err
	}
	cert.Leaf = leaf
	return cert, nil
}

// ---- DNSimple zone records (the existing internal/bot client is registrar/
// OAuth-only; the router needs record create/delete with a long-lived token) --

type dnsimpleDNS struct {
	base    string
	token   string
	account string // discovered via /v2/whoami on first use
}

// newDNSimpleDNS builds the default DNS-01 backend from an account API token,
// honouring the LCHAT_DNSIMPLE_BASE override (kept here so all env reads for
// this backend live in one place).
func newDNSimpleDNS(token string) *dnsimpleDNS {
	return &dnsimpleDNS{
		base:  envOr("LCHAT_DNSIMPLE_BASE", "https://api.dnsimple.com"),
		token: token,
	}
}

// present satisfies dnsProvider: create the challenge TXT, then hand back a
// closure that deletes it by the record id DNSimple returned. createTXT and
// deleteRecord are unchanged, so the DNSimple wire behaviour is byte-for-byte
// what it was before the seam existed.
func (d *dnsimpleDNS) present(ctx context.Context, zone, name, value string) (func(context.Context) error, error) {
	id, err := d.createTXT(ctx, zone, name, value)
	if err != nil {
		return nil, fmt.Errorf("dnsimple: create TXT: %v", err)
	}
	return func(cctx context.Context) error { return d.deleteRecord(cctx, zone, id) }, nil
}

func (d *dnsimpleDNS) do(ctx context.Context, method, path string, body, out any) error {
	var rdr io.Reader
	if body != nil {
		raw, err := json.Marshal(body)
		if err != nil {
			return err
		}
		rdr = bytes.NewReader(raw)
	}
	req, err := http.NewRequestWithContext(ctx, method, d.base+path, rdr)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+d.token)
	req.Header.Set("Accept", "application/json")
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return fmt.Errorf("%s %s: %s: %s", method, path, resp.Status, strings.TrimSpace(string(msg)))
	}
	if out != nil {
		return json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(out)
	}
	return nil
}

func (d *dnsimpleDNS) accountID(ctx context.Context) (string, error) {
	if d.account != "" {
		return d.account, nil
	}
	var who struct {
		Data struct {
			Account *struct {
				ID int64 `json:"id"`
			} `json:"account"`
		} `json:"data"`
	}
	if err := d.do(ctx, "GET", "/v2/whoami", nil, &who); err != nil {
		return "", err
	}
	if who.Data.Account == nil {
		return "", errors.New("whoami: token is not an account token")
	}
	d.account = fmt.Sprint(who.Data.Account.ID)
	return d.account, nil
}

func (d *dnsimpleDNS) createTXT(ctx context.Context, zone, name, content string) (int64, error) {
	acct, err := d.accountID(ctx)
	if err != nil {
		return 0, err
	}
	var created struct {
		Data struct {
			ID int64 `json:"id"`
		} `json:"data"`
	}
	err = d.do(ctx, "POST", "/v2/"+acct+"/zones/"+zone+"/records", map[string]any{
		"name": name, "type": "TXT", "content": content, "ttl": 60,
	}, &created)
	if err != nil {
		return 0, err
	}
	return created.Data.ID, nil
}

func (d *dnsimpleDNS) deleteRecord(ctx context.Context, zone string, id int64) error {
	acct, err := d.accountID(ctx)
	if err != nil {
		return err
	}
	return d.do(ctx, "DELETE", fmt.Sprintf("/v2/%s/zones/%s/records/%d", acct, zone, id), nil, nil)
}

// splitDomains parses the LCHAT_WILDCARD_DOMAINS comma list.
func splitDomains(s string) []string {
	var out []string
	for _, d := range strings.Split(s, ",") {
		if d = strings.ToLower(strings.TrimSpace(d)); d != "" {
			out = append(out, d)
		}
	}
	return out
}