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 (also resolves IPs for PTR check)
	ips, aRecordChecks := checkARecords(domain)
	checks = append(checks, aRecordChecks...)

	// 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)...)

	// PTR / rDNS match
	checks = append(checks, checkPTR(ips)...)

	return types.DNSResult{Checks: checks}
}

// checkARecords returns the resolved IPv4 addresses and the check results.
func checkARecords(domain string) ([]string, []types.CheckResult) {
	ips, err := net.LookupHost(domain)
	if err != nil {
		return nil, []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 nil, []types.CheckResult{{Status: types.StatusFail, Message: "no A record found"}}
	}
	return ipv4, []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
	}
	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

	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") {
				checks = append(checks, types.CheckResult{Status: types.StatusOK, Message: "DMARC record found"})
				return checks
			}
		}
	}
	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 {
	_, err := net.LookupTXT("_dnssec." + domain)
	if err == nil {
		return []types.CheckResult{{Status: types.StatusOK, Message: "DNSSEC appears enabled"}}
	}
	return []types.CheckResult{{Status: types.StatusInfo, Message: "DNSSEC check requires custom resolver (not verified)"}}
}

// checkPTR performs a reverse DNS lookup on the given IPs and checks if
// any returned hostname resolves back to one of the original IPs.
func checkPTR(ips []string) []types.CheckResult {
	if len(ips) == 0 {
		return nil
	}

	// Check up to the first 2 IPs to keep things fast.
	limit := 2
	if len(ips) < limit {
		limit = len(ips)
	}

	for i := 0; i < limit; i++ {
		ip := ips[i]
		names, err := net.LookupAddr(ip)
		if err != nil || len(names) == 0 {
			continue
		}

		for _, name := range names {
			name = strings.TrimSuffix(name, ".")
			resolved, err := net.LookupHost(name)
			if err != nil {
				continue
			}
			for _, resolvedIP := range resolved {
				if resolvedIP == ip {
					return []types.CheckResult{{Status: types.StatusOK, Message: fmt.Sprintf("rDNS matches (%s → %s)", ip, name)}}
				}
			}
		}
	}

	return []types.CheckResult{{Status: types.StatusInfo, Message: "no PTR record found (rDNS not configured)"}}
}