package acme

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"time"
)

// dnsResolver is the DNS resolver used for DNS-01 challenge validation.
// It defaults to the system resolver and can be replaced in tests.
var dnsResolver = net.DefaultResolver

// validateChallenge dispatches to the appropriate validator and updates
// challenge, authorization, and order state in the barrier.
func (h *Handler) validateChallenge(ctx context.Context, chall *Challenge, accountJWK []byte) {
	// Load the authorization to get the identifier (domain/IP).
	authz, err := h.loadAuthz(ctx, chall.AuthzID)
	if err != nil {
		h.logger.Error("acme: load authz for validation", "id", chall.AuthzID, "error", err)
		chall.Status = StatusInvalid
		chall.Error = &ProblemDetail{Type: ProblemServerInternal, Detail: "failed to load authorization"}
		_ = h.saveChallenge(ctx, chall)
		return
	}
	// Inject the identifier value into the context for validators.
	ctx = context.WithValue(ctx, ctxKeyDomain, authz.Identifier.Value)

	var validationErr error
	switch chall.Type {
	case ChallengeHTTP01:
		validationErr = validateHTTP01(ctx, chall, accountJWK)
	case ChallengeDNS01:
		validationErr = validateDNS01(ctx, chall, accountJWK)
	default:
		validationErr = fmt.Errorf("unknown challenge type: %s", chall.Type)
	}

	now := time.Now()
	if validationErr == nil {
		chall.Status = StatusValid
		chall.ValidatedAt = &now
		chall.Error = nil
	} else {
		h.logger.Info("acme: challenge validation failed",
			"id", chall.ID, "type", chall.Type, "error", validationErr)
		chall.Status = StatusInvalid
		chall.Error = &ProblemDetail{
			Type:   ProblemConnection,
			Detail: validationErr.Error(),
		}
	}

	if err := h.saveChallenge(ctx, chall); err != nil {
		h.logger.Error("acme: save challenge after validation", "error", err)
		return
	}

	// Update the parent authorization.
	h.updateAuthzStatus(ctx, chall.AuthzID)
}

// updateAuthzStatus recomputes an authorization's status from its challenges,
// then propagates any state change to the parent order.
func (h *Handler) updateAuthzStatus(ctx context.Context, authzID string) {
	authz, err := h.loadAuthz(ctx, authzID)
	if err != nil {
		h.logger.Error("acme: load authz for status update", "error", err)
		return
	}

	// An authz is valid if any single challenge is valid.
	// An authz is invalid if all challenges are invalid.
	anyValid := false
	allInvalid := true
	for _, challID := range authz.ChallengeIDs {
		chall, err := h.loadChallenge(ctx, challID)
		if err != nil {
			continue
		}
		if chall.Status == StatusValid {
			anyValid = true
			allInvalid = false
			break
		}
		if chall.Status != StatusInvalid {
			allInvalid = false
		}
	}

	prevStatus := authz.Status
	if anyValid {
		authz.Status = StatusValid
	} else if allInvalid {
		authz.Status = StatusInvalid
	}

	if authz.Status != prevStatus {
		data, _ := json.Marshal(authz)
		if err := h.barrier.Put(ctx, h.barrierPrefix()+"authz/"+authzID+".json", data); err != nil {
			h.logger.Error("acme: save authz", "error", err)
			return
		}
		// Propagate to orders that reference this authz.
		h.updateOrdersForAuthz(ctx, authzID)
	}
}

// updateOrdersForAuthz scans all orders for one that references authzID
// and updates the order status if all authorizations are now valid.
func (h *Handler) updateOrdersForAuthz(ctx context.Context, authzID string) {
	paths, err := h.barrier.List(ctx, h.barrierPrefix()+"orders/")
	if err != nil {
		return
	}
	for _, p := range paths {
		data, err := h.barrier.Get(ctx, h.barrierPrefix()+"orders/"+p)
		if err != nil || data == nil {
			continue
		}
		var order Order
		if err := json.Unmarshal(data, &order); err != nil {
			continue
		}
		if order.Status != StatusPending {
			continue
		}
		for _, id := range order.AuthzIDs {
			if id != authzID {
				continue
			}
			// This order references the updated authz; check all authzs.
			h.maybeAdvanceOrder(ctx, &order)
			break
		}
	}
}

// maybeAdvanceOrder checks whether all authorizations for an order are valid
// and transitions the order to "ready" if so.
func (h *Handler) maybeAdvanceOrder(ctx context.Context, order *Order) {
	allValid := true
	for _, authzID := range order.AuthzIDs {
		authz, err := h.loadAuthz(ctx, authzID)
		if err != nil || authz.Status != StatusValid {
			allValid = false
			break
		}
	}
	if !allValid {
		return
	}
	order.Status = StatusReady
	data, _ := json.Marshal(order)
	if err := h.barrier.Put(ctx, h.barrierPrefix()+"orders/"+order.ID+".json", data); err != nil {
		h.logger.Error("acme: advance order to ready", "error", err)
	}
}

// validateHTTP01 performs HTTP-01 challenge validation (RFC 8555 §8.3).
// It fetches http://{domain}/.well-known/acme-challenge/{token} and verifies
// the response matches the key authorization.
func validateHTTP01(ctx context.Context, chall *Challenge, accountJWK []byte) error {
	authz, err := KeyAuthorization(chall.Token, accountJWK)
	if err != nil {
		return fmt.Errorf("compute key authorization: %w", err)
	}

	// Load the authorization to get the identifier.
	// The domain comes from the parent authorization; we pass accountJWK here
	// but need the identifier. It's embedded in the challenge's AuthzID so the
	// caller (validateChallenge) must load the authz to get the domain.
	// Since we don't have the domain directly in Challenge, we embed it during
	// challenge creation. For now: the caller passes accountJWK, and we use
	// a context value to carry the domain through validateHTTP01.
	domain, ok := ctx.Value(ctxKeyDomain).(string)
	if !ok || domain == "" {
		return errors.New("domain not in context")
	}

	url := "http://" + domain + "/.well-known/acme-challenge/" + chall.Token
	client := &http.Client{
		Timeout: 10 * time.Second,
		// Follow at most 10 redirects per RFC 8555 §8.3.
		CheckRedirect: func(_ *http.Request, via []*http.Request) error {
			if len(via) >= 10 {
				return errors.New("too many redirects")
			}
			return nil
		},
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return fmt.Errorf("build HTTP-01 request: %w", err)
	}

	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("HTTP-01 fetch failed: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP-01: unexpected status %d", resp.StatusCode)
	}

	body, err := io.ReadAll(io.LimitReader(resp.Body, 512))
	if err != nil {
		return fmt.Errorf("HTTP-01: read body: %w", err)
	}

	// RFC 8555 §8.3: response body should be the key authorization, optionally
	// followed by whitespace.
	got := string(body)
	for len(got) > 0 && (got[len(got)-1] == '\n' || got[len(got)-1] == '\r' || got[len(got)-1] == ' ') {
		got = got[:len(got)-1]
	}
	if got != authz {
		return fmt.Errorf("HTTP-01: key authorization mismatch (got %q, want %q)", got, authz)
	}
	return nil
}

// validateDNS01 performs DNS-01 challenge validation (RFC 8555 §8.4).
// It looks up _acme-challenge.{domain} TXT records and checks for the
// base64url(SHA-256(keyAuthorization)) value.
func validateDNS01(ctx context.Context, chall *Challenge, accountJWK []byte) error {
	keyAuth, err := KeyAuthorization(chall.Token, accountJWK)
	if err != nil {
		return fmt.Errorf("compute key authorization: %w", err)
	}
	// DNS-01 TXT record value is base64url(SHA-256(keyAuthorization)).
	digest := sha256.Sum256([]byte(keyAuth))
	expected := base64.RawURLEncoding.EncodeToString(digest[:])

	domain, ok := ctx.Value(ctxKeyDomain).(string)
	if !ok || domain == "" {
		return errors.New("domain not in context")
	}

	// Strip trailing dot if present; add _acme-challenge prefix.
	domain = "_acme-challenge." + domain

	txts, err := dnsResolver.LookupTXT(ctx, domain)
	if err != nil {
		return fmt.Errorf("DNS-01: TXT lookup for %s failed: %w", domain, err)
	}

	for _, txt := range txts {
		if txt == expected {
			return nil
		}
	}
	return fmt.Errorf("DNS-01: no matching TXT record for %s (expected %s)", domain, expected)
}

// ctxKeyDomain is the context key for passing the domain to validators.
type ctxDomainKey struct{}

var ctxKeyDomain = ctxDomainKey{}

// pemToDER decodes the first PEM block from a PEM string and returns its DER bytes.
func pemToDER(pemStr string) ([]byte, error) {
	block, _ := pem.Decode([]byte(pemStr))
	if block == nil {
		return nil, errors.New("no PEM block")
	}
	return block.Bytes, nil
}
