package mailstore

import (
	"context"
	"net"
	"strings"
)

// CheckSPF is a minimal, single-level SPF check (v=spf1 ip4:/a/mx/include:, no
// recursive include/redirect, no macro expansion) against the sender domain's TXT
// record — one signal feeding the spam heuristic in spam.go, not an authoritative
// pass/fail gate.
// ponytail: not full RFC 7208 (no multi-level includes, no redirect, no macros) —
// good enough as a signal, revisit if a real sender's SPF record depends on it.
func CheckSPF(mailFrom, peerIP string) bool {
	domain := domainOf(mailFrom)
	if domain == "" {
		return true
	}
	ip := net.ParseIP(peerIP)
	if ip == nil {
		return true
	}
	record, ok := lookupSPFRecord(domain)
	if !ok {
		return true // no SPF record published: neutral, not a penalty
	}
	return evalSPF(record, ip, domain, 0)
}

func domainOf(address string) string {
	i := strings.LastIndex(address, "@")
	if i < 0 {
		return ""
	}
	return strings.ToLower(address[i+1:])
}

func lookupSPFRecord(domain string) (string, bool) {
	txts, err := net.DefaultResolver.LookupTXT(context.Background(), domain)
	if err != nil {
		return "", false
	}
	for _, t := range txts {
		if strings.HasPrefix(strings.ToLower(t), "v=spf1") {
			return t, true
		}
	}
	return "", false
}

// evalSPF walks mechanisms left to right; depth caps includes at one level.
func evalSPF(record string, ip net.IP, domain string, depth int) bool {
	fields := strings.Fields(record)
	for _, f := range fields[1:] { // skip "v=spf1"
		qualifier := byte('+')
		mech := f
		if len(f) > 0 && strings.ContainsRune("+-~?", rune(f[0])) {
			qualifier = f[0]
			mech = f[1:]
		}
		switch {
		case mech == "all":
			return qualifier != '-'
		case strings.HasPrefix(mech, "ip4:"):
			if matchIP4(mech[4:], ip) {
				return qualifier != '-'
			}
		case mech == "a":
			if matchA(domain, ip) {
				return qualifier != '-'
			}
		case strings.HasPrefix(mech, "a:"):
			if matchA(mech[2:], ip) {
				return qualifier != '-'
			}
		case mech == "mx":
			if matchMX(domain, ip) {
				return qualifier != '-'
			}
		case strings.HasPrefix(mech, "mx:"):
			if matchMX(mech[3:], ip) {
				return qualifier != '-'
			}
		case strings.HasPrefix(mech, "include:") && depth == 0:
			sub, ok := lookupSPFRecord(mech[len("include:"):])
			if ok && evalSPF(sub, ip, mech[len("include:"):], depth+1) {
				return true
			}
		}
	}
	return true // no matching mechanism and no explicit "all": neutral
}

func matchIP4(cidr string, ip net.IP) bool {
	if !strings.Contains(cidr, "/") {
		cidr += "/32"
	}
	_, network, err := net.ParseCIDR(cidr)
	if err != nil {
		return false
	}
	return network.Contains(ip)
}

func matchA(host string, ip net.IP) bool {
	ips, err := net.DefaultResolver.LookupIP(context.Background(), "ip4", host)
	if err != nil {
		return false
	}
	for _, a := range ips {
		if a.Equal(ip) {
			return true
		}
	}
	return false
}

func matchMX(domain string, ip net.IP) bool {
	mxs, err := net.LookupMX(domain)
	if err != nil {
		return false
	}
	for _, mx := range mxs {
		if matchA(strings.TrimSuffix(mx.Host, "."), ip) {
			return true
		}
	}
	return false
}