package dns

import (
	"fmt"
	"net"
	"strings"

	"arcline-audit/internal/types"
)

// Run performs all DNS checks for the given domain.
func Run(domain string) types.DNSResult {
	var checks []types.CheckResult

	// A records
	checks = append(checks, checkARecords(domain)...)

	// AAAA records
	checks = append(checks, checkAAAARecords(domain)...)

	// MX records
	checks = append(checks, checkMXRecords(domain)...)

	// TXT records (SPF, DKIM, DMARC)
	checks = append(checks, checkTXTRecords(domain)...)

	// DNSSEC
	checks = append(checks, checkDNSSEC(domain)...)

	return types.DNSResult{Checks: checks}
}

func checkARecords(domain string) []types.CheckResult {
	ips, err := net.LookupHost(domain)
	if err != nil {
		return []types.CheckResult{{Status: types.StatusFail, Message: fmt.Sprintf("no A record: %v", err)}}
	}
	var ipv4 []string
	for _, ip := range ips {
		if parsed := net.ParseIP(ip); parsed != nil && parsed.To4() != nil {
			ipv4 = append(ipv4, ip)
		}
	}
	if len(ipv4) == 0 {
		return []types.CheckResult{{Status: types.StatusFail, Message: "no A record found"}}
	}
	return []types.CheckResult{{Status: types.StatusOK, Message: fmt.Sprintf("A record: %s", strings.Join(ipv4, ", "))}}
}

func checkAAAARecords(domain string) []types.CheckResult {
	ips, err := net.LookupHost(domain)
	if err != nil {
		return nil
	}
	var ipv6 []string
	for _, ip := range ips {
		if parsed := net.ParseIP(ip); parsed != nil && parsed.To4() == nil {
			ipv6 = append(ipv6, ip)
		}
	}
	if len(ipv6) == 0 {
		return nil // Not a warning; many sites don't have IPv6
	}
	return []types.CheckResult{{Status: types.StatusOK, Message: fmt.Sprintf("AAAA record: %s", strings.Join(ipv6, ", "))}}
}

func checkMXRecords(domain string) []types.CheckResult {
	mxs, err := net.LookupMX(domain)
	if err != nil || len(mxs) == 0 {
		return []types.CheckResult{{Status: types.StatusWarn, Message: "no MX records found"}}
	}
	return []types.CheckResult{{Status: types.StatusOK, Message: fmt.Sprintf("MX records present (%d)", len(mxs))}}
}

func checkTXTRecords(domain string) []types.CheckResult {
	txts, err := net.LookupTXT(domain)
	if err != nil {
		return []types.CheckResult{
			{Status: types.StatusWarn, Message: "no TXT records found"},
		}
	}

	var checks []types.CheckResult
	hasSPF := false
	hasDMARC := false

	for _, txt := range txts {
		if strings.HasPrefix(txt, "v=spf1") {
			hasSPF = true
		}
	}

	if hasSPF {
		checks = append(checks, types.CheckResult{Status: types.StatusOK, Message: "SPF record found"})
	} else {
		checks = append(checks, types.CheckResult{Status: types.StatusWarn, Message: "no SPF record"})
	}

	// DKIM is checked via selector lookup
	dkimFound := checkDKIM(domain)
	if dkimFound {
		checks = append(checks, types.CheckResult{Status: types.StatusOK, Message: "DKIM record found (default._domainkey)"})
	} else {
		checks = append(checks, types.CheckResult{Status: types.StatusInfo, Message: "no DKIM record (default._domainkey)"})
	}

	// DMARC is on _dmarc subdomain
	dmarcTxts, err := net.LookupTXT("_dmarc." + domain)
	if err == nil {
		for _, txt := range dmarcTxts {
			if strings.HasPrefix(txt, "v=DMARC1") {
				hasDMARC = true
				break
			}
		}
	}
	if hasDMARC {
		checks = append(checks, types.CheckResult{Status: types.StatusOK, Message: "DMARC record found"})
	} else {
		checks = append(checks, types.CheckResult{Status: types.StatusWarn, Message: "no DMARC record"})
	}

	return checks
}

func checkDKIM(domain string) bool {
	txts, err := net.LookupTXT("default._domainkey." + domain)
	if err != nil {
		return false
	}
	for _, txt := range txts {
		if strings.Contains(txt, "v=DKIM1") {
			return true
		}
	}
	return false
}

func checkDNSSEC(domain string) []types.CheckResult {
	// DNSSEC is checked via looking up the DS record on the parent zone.
	// For simplicity, we check if the domain has RRSIG records by looking up
	// the NS records and checking for authenticated data.
	// A true DNSSEC check requires a validating resolver. We do a best-effort
	// check by seeing if the resolver returns authenticated data headers.
	// As a simple heuristic, we check for DNSKEY records.
	_, err := net.LookupTXT("_dnssec." + domain)
	if err == nil {
		return []types.CheckResult{{Status: types.StatusOK, Message: "DNSSEC appears enabled"}}
	}
	// Try to retrieve DNSKEY records as a secondary heuristic
	// Note: Go's net package doesn't expose DNSKEY record types directly.
	// A full DNSSEC check would require a custom DNS resolver.
	return []types.CheckResult{{Status: types.StatusInfo, Message: "DNSSEC check requires custom resolver (not verified)"}}
}