package server

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

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

// IssuedCert is a certificate produced for a tunnel's CSR.
type IssuedCert struct {
	Certificate string // leaf PEM
	Chain       string // issuer chain PEM
	Expiry      time.Time
}

// Issuer turns a tunnel's CSR into a certificate. The local CA is instant;
// ACME takes a DNS-01 round trip, so callers run those asynchronously.
type Issuer interface {
	Instant() bool
	Issue(ctx context.Context, csrPEM, hostname string) (IssuedCert, error)
}

// ---- local CA issuer (development and self-hosting without ACME).

type localIssuer struct{ store *Store }

// NewLocalIssuer signs tunnel CSRs with the store's built-in CA.
func NewLocalIssuer(store *Store) Issuer { return &localIssuer{store: store} }

func (l *localIssuer) Instant() bool { return true }

func (l *localIssuer) Issue(_ context.Context, csrPEM, hostname string) (IssuedCert, error) {
	ca, err := l.store.ca()
	if err != nil {
		return IssuedCert{}, err
	}
	certPEM, chainPEM, expiry, err := ca.signCSR(csrPEM, hostname)
	if err != nil {
		return IssuedCert{}, err
	}
	return IssuedCert{Certificate: string(certPEM), Chain: string(chainPEM), Expiry: expiry}, nil
}

// ---- ACME issuer (Let's Encrypt, ZeroSSL, Pebble) with DNS-01 challenges.

type ACMEOptions struct {
	DirectoryURL string
	Email        string
	CABundleFile string // extra roots for the ACME server's TLS (Pebble)
	DNS          DNSProvider
	// PropagationTimeout waits until public resolvers see the challenge
	// records; 0 skips the check (used with pebble-challtestsrv).
	PropagationTimeout time.Duration
}

type acmeIssuer struct {
	store *Store
	opts  ACMEOptions
	sem   chan struct{} // bounds concurrent ACME orders

	accountMu sync.Mutex // serializes first-use account registration
}

// NewACMEIssuer returns an issuer that orders certificates over ACME and
// answers DNS-01 challenges through the given DNS provider.
func NewACMEIssuer(store *Store, opts ACMEOptions) (Issuer, error) {
	if opts.DirectoryURL == "" {
		return nil, errors.New("acme: directory URL is required")
	}
	if opts.DNS == nil {
		return nil, errors.New("acme: a DNS provider is required")
	}
	return &acmeIssuer{store: store, opts: opts, sem: make(chan struct{}, 2)}, nil
}

func (a *acmeIssuer) Instant() bool { return false }

// client returns an ACME client with the persisted account (registering one
// on first use).
func (a *acmeIssuer) client(ctx context.Context) (*acme.Client, error) {
	a.accountMu.Lock()
	defer a.accountMu.Unlock()

	keyPEM, accountURL, directory := a.store.acmeAccount()
	storedKey, storedURL, storedDirectory := keyPEM, accountURL, directory
	if directory != a.opts.DirectoryURL {
		// The persisted account belongs to another ACME server (for example
		// staging); register it with this directory instead.
		accountURL = ""
	}

	var key *ecdsa.PrivateKey
	if keyPEM != "" {
		block, _ := pem.Decode([]byte(keyPEM))
		if block == nil {
			return nil, errors.New("acme: unreadable account key")
		}
		parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("acme: account key: %w", err)
		}
		ec, ok := parsed.(*ecdsa.PrivateKey)
		if !ok {
			return nil, errors.New("acme: account key is not ECDSA")
		}
		key = ec
	} else {
		generated, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
		if err != nil {
			return nil, err
		}
		key = generated
		der, err := x509.MarshalPKCS8PrivateKey(key)
		if err != nil {
			return nil, err
		}
		keyPEM = string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der}))
	}

	httpClient := &http.Client{Timeout: 60 * time.Second}
	if a.opts.CABundleFile != "" {
		pool, err := x509.SystemCertPool()
		if err != nil {
			pool = x509.NewCertPool()
		}
		bundle, err := os.ReadFile(a.opts.CABundleFile)
		if err != nil {
			return nil, fmt.Errorf("acme: CA bundle: %w", err)
		}
		if !pool.AppendCertsFromPEM(bundle) {
			return nil, errors.New("acme: CA bundle has no certificates")
		}
		transport := http.DefaultTransport.(*http.Transport).Clone()
		transport.TLSClientConfig = &tls.Config{RootCAs: pool}
		httpClient.Transport = transport
	}

	client := &acme.Client{Key: key, DirectoryURL: a.opts.DirectoryURL, HTTPClient: httpClient}
	if accountURL != "" {
		// Reuse the persisted account without a GetAccount round trip.
		client.KID = acme.KeyID(accountURL)
	}
	if accountURL == "" {
		account := &acme.Account{}
		if a.opts.Email != "" {
			account.Contact = []string{"mailto:" + a.opts.Email}
		}
		registered, err := client.Register(ctx, account, acme.AcceptTOS)
		switch {
		case errors.Is(err, acme.ErrAccountAlreadyExists):
			// The key is already registered with this CA; Register cached the
			// account URL in client.KID.
			accountURL = string(client.KID)
		case err != nil:
			return nil, fmt.Errorf("acme: register account: %w", err)
		default:
			accountURL = registered.URI
		}
	}
	if keyPEM != storedKey || accountURL != storedURL || a.opts.DirectoryURL != storedDirectory {
		if err := a.store.setACMEAccount(keyPEM, accountURL, a.opts.DirectoryURL); err != nil {
			return nil, err
		}
	}
	return client, nil
}

// Issue orders a certificate for the tunnel's names (hostname and, when the
// CSR asks for it, its wildcard) using one DNS-01 challenge record shared by
// both, then returns the certificate for the client's key.
func (a *acmeIssuer) Issue(ctx context.Context, csrPEM, hostname string) (IssuedCert, error) {
	csr, err := parseCSR(csrPEM, hostname)
	if err != nil {
		return IssuedCert{}, err
	}

	select {
	case a.sem <- struct{}{}:
	case <-ctx.Done():
		return IssuedCert{}, ctx.Err()
	}
	defer func() { <-a.sem }()

	client, err := a.client(ctx)
	if err != nil {
		return IssuedCert{}, err
	}

	ids := []acme.AuthzID{{Type: "dns", Value: hostname}}
	for _, name := range csr.DNSNames {
		if name == "*."+hostname {
			ids = append(ids, acme.AuthzID{Type: "dns", Value: name})
		}
	}
	order, err := client.AuthorizeOrder(ctx, ids)
	if err != nil {
		return IssuedCert{}, fmt.Errorf("acme: authorize order: %w", err)
	}
	if order.URI == "" || order.FinalizeURL == "" {
		return IssuedCert{}, fmt.Errorf("acme: order is missing its URI or finalize URL (status %q)", order.Status)
	}

	record := challengeRecordName(hostname)
	var values []string
	created := false
	defer func() {
		if !created {
			return
		}
		// Best-effort cleanup, even when the attempt failed or the context is
		// done; a leftover record would block the next attempt.
		cleanCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		if err := a.opts.DNS.DeleteTXT(cleanCtx, record); err != nil {
			log.Printf("acme: cleanup TXT %s: %v", record, err)
		}
	}()
	for _, authzURL := range order.AuthzURLs {
		authz, err := client.GetAuthorization(ctx, authzURL)
		if err != nil {
			return IssuedCert{}, fmt.Errorf("acme: read authorization: %w", err)
		}
		if authz.Status == acme.StatusValid {
			continue
		}
		if authz.Status != acme.StatusPending {
			return IssuedCert{}, authzFailure(authz)
		}
		challenge := dns01Challenge(authz)
		if challenge == nil {
			return IssuedCert{}, fmt.Errorf("acme: no dns-01 challenge for %s", authz.Identifier.Value)
		}
		value, err := client.DNS01ChallengeRecord(challenge.Token)
		if err != nil {
			return IssuedCert{}, fmt.Errorf("acme: challenge record: %w", err)
		}
		if err := a.opts.DNS.CreateTXT(ctx, record, value); err != nil {
			return IssuedCert{}, fmt.Errorf("dns: create %s: %w", record, err)
		}
		created = true
		values = append(values, value)
	}

	if a.opts.PropagationTimeout > 0 {
		if err := waitForTXT(ctx, record, values, a.opts.PropagationTimeout); err != nil {
			return IssuedCert{}, err
		}
	}

	for _, authzURL := range order.AuthzURLs {
		authz, err := client.GetAuthorization(ctx, authzURL)
		if err != nil {
			return IssuedCert{}, fmt.Errorf("acme: read authorization: %w", err)
		}
		if authz.Status == acme.StatusValid {
			continue
		}
		if authz.Status != acme.StatusPending {
			return IssuedCert{}, authzFailure(authz)
		}
		challenge := dns01Challenge(authz)
		if challenge == nil {
			continue
		}
		if _, err := client.Accept(ctx, challenge); err != nil {
			return IssuedCert{}, fmt.Errorf("acme: accept challenge: %w", err)
		}
		if _, err := client.WaitAuthorization(ctx, authzURL); err != nil {
			return IssuedCert{}, fmt.Errorf("acme: authorization: %w", err)
		}
	}

	// All authorizations are satisfied; wait for the CA to mark the order
	// ready before submitting the CSR (Let's Encrypt rejects finalization
	// while the order is still pending).
	order, err = client.WaitOrder(ctx, order.URI)
	if err != nil {
		return IssuedCert{}, fmt.Errorf("acme: wait order: %w", err)
	}
	var der [][]byte
	if order.Status == acme.StatusValid && order.CertURL != "" {
		der, err = client.FetchCert(ctx, order.CertURL, true)
	} else {
		der, _, err = client.CreateCertFromOrder(ctx, order, csr.Raw, true)
	}
	if err != nil {
		return IssuedCert{}, fmt.Errorf("acme: finalize order: %w", err)
	}
	if len(der) == 0 {
		return IssuedCert{}, errors.New("acme: empty certificate chain")
	}
	leaf, err := x509.ParseCertificate(der[0])
	if err != nil {
		return IssuedCert{}, fmt.Errorf("acme: parse certificate: %w", err)
	}
	var certPEM, chainPEM strings.Builder
	for i, certDER := range der {
		block := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
		if i == 0 {
			certPEM.Write(block)
		} else {
			chainPEM.Write(block)
		}
	}
	return IssuedCert{Certificate: certPEM.String(), Chain: chainPEM.String(), Expiry: leaf.NotAfter}, nil
}

func dns01Challenge(authz *acme.Authorization) *acme.Challenge {
	for _, challenge := range authz.Challenges {
		if challenge.Type == "dns-01" {
			return challenge
		}
	}
	return nil
}

// authzFailure explains why an authorization cannot be used, including the
// CA's validation error when it reported one.
func authzFailure(authz *acme.Authorization) error {
	for _, challenge := range authz.Challenges {
		if challenge.Error != nil {
			return fmt.Errorf("acme: authorization for %s is %s: %w", authz.Identifier.Value, authz.Status, challenge.Error)
		}
	}
	return fmt.Errorf("acme: authorization for %s is %s", authz.Identifier.Value, authz.Status)
}

// waitForTXT polls until every challenge value is visible at the record name.
// It tries the system resolver first, then public resolvers directly, so a
// cached negative answer on the container's resolver cannot hide a record
// that the CA will see.
func waitForTXT(ctx context.Context, record string, values []string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	resolvers := []*net.Resolver{net.DefaultResolver, resolverAt("1.1.1.1:53"), resolverAt("8.8.8.8:53")}
	for {
		for _, resolver := range resolvers {
			queryCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
			seen, err := resolver.LookupTXT(queryCtx, record)
			cancel()
			if err != nil {
				continue
			}
			present := 0
			for _, want := range values {
				for _, got := range seen {
					if got == want {
						present++
						break
					}
				}
			}
			if present == len(values) {
				return nil
			}
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("dns: TXT record %s did not propagate within %s", record, timeout)
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
}

// resolverAt queries one DNS server directly instead of the system resolver.
func resolverAt(address string) *net.Resolver {
	return &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, _ string) (net.Conn, error) {
			dialer := net.Dialer{Timeout: 4 * time.Second}
			return dialer.DialContext(ctx, network, address)
		},
	}
}