// Package mtasts implements RFC 8461 MTA-STS policy discovery for outbound
// SMTP delivery — a domain publishes a DNS TXT record plus an HTTPS-hosted
// policy document declaring which MX hosts must be used and whether TLS is
// mandatory. Unlike DANE, this doesn't depend on DNSSEC: the policy fetch's
// own TLS certificate (normal CA/PKI, already handled by net/http) is the
// trust anchor, per the RFC.
package mtasts

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

// Policy is a parsed MTA-STS policy document.
type Policy struct {
	ID         string // from the _mta-sts TXT record, not the policy body
	Mode       string // "enforce" | "testing" | "none"
	MXPatterns []string
	MaxAge     time.Duration
}

// Discover fetches domain's MTA-STS policy. Returns (nil, nil) — not an
// error — if the domain has no _mta-sts TXT record at all, since that's the
// normal "this domain doesn't use MTA-STS" case.
func Discover(ctx context.Context, domain string) (*Policy, error) {
	// Same net.DefaultResolver.LookupTXT convention already used by
	// internal/pipeline's SPF/DMARC checks — no AD-flag need here, so no
	// reason to use the hand-rolled dnsutil client for this lookup.
	txts, err := net.DefaultResolver.LookupTXT(ctx, "_mta-sts."+domain)
	if err != nil || len(txts) == 0 {
		return nil, nil
	}

	var policyID string
	for _, txt := range txts {
		if !strings.HasPrefix(txt, "v=STSv1") {
			continue
		}
		for _, part := range strings.Split(txt, ";") {
			part = strings.TrimSpace(part)
			if id, ok := strings.CutPrefix(part, "id="); ok {
				policyID = id
			}
		}
	}
	if policyID == "" {
		return nil, nil
	}

	url := "https://mta-sts." + domain + "/.well-known/mta-sts.txt"
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("fetching mta-sts policy for %s: %w", domain, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("mta-sts policy fetch for %s returned status %d", domain, resp.StatusCode)
	}

	// Policies are meant to be small (a handful of mx lines) — cap the read
	// against a hostile or misbehaving server rather than trusting Content-Length.
	body, err := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
	if err != nil {
		return nil, fmt.Errorf("reading mta-sts policy body: %w", err)
	}

	policy, err := parsePolicy(string(body))
	if err != nil {
		return nil, fmt.Errorf("parsing mta-sts policy for %s: %w", domain, err)
	}
	policy.ID = policyID
	return policy, nil
}

func parsePolicy(body string) (*Policy, error) {
	p := &Policy{}
	var maxAgeSeconds int
	for _, line := range strings.Split(body, "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		key, value, ok := strings.Cut(line, ":")
		if !ok {
			continue
		}
		key = strings.TrimSpace(key)
		value = strings.TrimSpace(value)
		switch key {
		case "version":
			if value != "STSv1" {
				return nil, fmt.Errorf("unsupported policy version %q", value)
			}
		case "mode":
			p.Mode = value
		case "mx":
			p.MXPatterns = append(p.MXPatterns, value)
		case "max_age":
			if n, err := strconv.Atoi(value); err == nil {
				maxAgeSeconds = n
			}
		}
	}
	if p.Mode == "" {
		return nil, fmt.Errorf("policy missing required 'mode' field")
	}
	p.MaxAge = time.Duration(maxAgeSeconds) * time.Second
	return p, nil
}

// Matches reports whether host satisfies pattern, per RFC 8461 §4.1's
// one-label wildcard rule: "*.example.com" matches "mail.example.com" but
// not "example.com" itself or "a.mail.example.com".
func Matches(pattern, host string) bool {
	pattern = strings.TrimSuffix(strings.ToLower(pattern), ".")
	host = strings.TrimSuffix(strings.ToLower(host), ".")

	suffix, isWildcard := strings.CutPrefix(pattern, "*.")
	if !isWildcard {
		return pattern == host
	}
	rest, ok := strings.CutSuffix(host, "."+suffix)
	return ok && rest != "" && !strings.Contains(rest, ".")
}