package router

import (
	"context"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"

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

var (
	certDir   = "/data/certs"
	domainsMu sync.RWMutex
	domains   = map[string]bool{}
	certMgr   *autocert.Manager
)

// AllowDomain validates DNS and adds a domain to the TLS whitelist.
// Returns an error if the domain doesn't resolve.
func AllowDomain(host string) error {
	host = strings.ToLower(strings.TrimSpace(host))
	if host == "" {
		return fmt.Errorf("empty hostname")
	}
	_, err := net.LookupHost(host)
	if err != nil {
		return fmt.Errorf("%s: DNS not found", host)
	}
	domainsMu.Lock()
	domains[host] = true
	domainsMu.Unlock()
	return nil
}

// AllowDomainUnchecked adds a domain without DNS validation.
// Use for system domains loaded from config/env at startup.
func AllowDomainUnchecked(host string) {
	domainsMu.Lock()
	domains[strings.ToLower(host)] = true
	domainsMu.Unlock()
}

// DisallowDomain removes a domain from the TLS whitelist.
func DisallowDomain(host string) {
	domainsMu.Lock()
	delete(domains, strings.ToLower(host))
	domainsMu.Unlock()
}

// SetCertDir sets the directory for autocert certificate storage.
// Must be called before Listen.
func SetCertDir(dir string) {
	certDir = dir
}

func hostPolicy(_ context.Context, host string) error {
	domainsMu.RLock()
	ok := domains[host]
	domainsMu.RUnlock()
	if !ok {
		return fmt.Errorf("router: %q not in allowed domains", host)
	}
	return nil
}

func ensureCertManager() {
	if certMgr != nil {
		return
	}
	os.MkdirAll(certDir, 0700)
	certMgr = &autocert.Manager{
		Prompt:     autocert.AcceptTOS,
		Cache:      autocert.DirCache(certDir),
		HostPolicy: hostPolicy,
	}
}