package blocklist

import (
	"bufio"
	"bytes"
	"fmt"
	"net/netip"
	"strconv"
	"strings"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/zonefile"
)

// Format identifies a list's syntax.
type Format string

const (
	// FormatAuto sniffs the content.
	FormatAuto Format = "auto"
	// FormatHosts is "IP domain [domain...]" (also accepts "0 domain").
	FormatHosts Format = "hosts"
	// FormatAdblock is the Adblock Plus / EasyList syntax.
	FormatAdblock Format = "adblock"
	// FormatPlain is one domain per line.
	FormatPlain Format = "plain"
	// FormatRPZ is an RFC 1035 response policy zone.
	FormatRPZ Format = "rpz"
)

// ParseOptions tunes the parsers.
type ParseOptions struct {
	// Subdomains makes a bare domain also cover its subdomains. Defaults to
	// true; hosts files rely on it, and precise lists can opt out.
	Subdomains bool
	// DefaultTTL is applied to rules that carry none.
	DefaultTTL uint32
	// Categories are attached to every rule from this source.
	Categories []string
}

// ParseStats reports what a parser saw, for the UI and for tests.
type ParseStats struct {
	Rules      int      `json:"rules"`
	Comments   int      `json:"comments"`
	Skipped    int      `json:"skipped"`
	Malformed  int      `json:"malformed"`
	Duplicates int      `json:"duplicates"`
	Bytes      int64    `json:"bytes"`
	DurationMS float64  `json:"duration_ms"`
	Format     Format   `json:"format"`
	Sample     []string `json:"sample,omitempty"`
	Notices    []string `json:"notices,omitempty"`
}

// DetectFormat sniffs a list's syntax from its first meaningful lines.
func DetectFormat(data []byte) Format {
	sc := bufio.NewScanner(bytes.NewReader(stripBOM(data)))
	sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	sawRecord := false
	for i := 0; sc.Scan() && i < 200; i++ {
		line := strings.TrimSpace(sc.Text())
		if line == "" || strings.HasPrefix(line, "#") || strings.HasPrefix(line, ";") {
			continue
		}
		if strings.HasPrefix(line, "[") {
			// "[Adblock Plus 2.0]" is a header, keep looking.
			continue
		}
		if strings.HasPrefix(line, "$TTL") || strings.HasPrefix(line, "$ORIGIN") ||
			strings.HasPrefix(line, "$INCLUDE") {
			return FormatRPZ
		}
		if sawRecord {
			return FormatAdblock
		}
		fields := strings.Fields(line)
		if len(fields) >= 2 && isAddr(fields[0]) {
			return FormatHosts
		}
		if strings.HasPrefix(line, "||") || strings.HasPrefix(line, "|") ||
			strings.HasPrefix(line, "@@") || strings.Contains(line, "##") ||
			strings.Contains(line, "#@#") || strings.Contains(line, "#?#") {
			return FormatAdblock
		}
		if len(fields) >= 3 && len(fields[1]) > 1 && allDigitsOrName(fields[1]) {
			// "name 60 IN A 1.2.3.4" style.
			return FormatRPZ
		}
		sawRecord = true
	}
	return FormatPlain
}

func allDigitsOrName(s string) bool {
	if _, err := strconv.Atoi(s); err == nil {
		return true
	}
	switch strings.ToUpper(s) {
	case "IN", "CH", "HS":
		return true
	}
	return false
}

// utf8BOM is the byte-order mark a UTF-8 file may start with. Editors on
// Windows add it, and it would otherwise glue itself to the first rule's
// owner name, silently dropping the first entry of the list.
var utf8BOM = []byte{0xef, 0xbb, 0xbf}

// stripBOM drops a leading UTF-8 byte-order mark. It is applied once, at the
// front door of the parse path, so every parser sees the same bytes.
func stripBOM(data []byte) []byte {
	return bytes.TrimPrefix(data, utf8BOM)
}

// Parse converts list data into rules using the given format (FormatAuto
// sniffs). It never returns partial garbage: on a hard error the rules
// accumulated so far are returned alongside it so the caller can decide
// whether to keep the last good set or adopt the partial one.
func Parse(format Format, data []byte, opt ParseOptions) ([]Rule, ParseStats, error) {
	data = stripBOM(data)
	if opt.DefaultTTL == 0 {
		opt.DefaultTTL = 60
	}
	if format == "" || format == FormatAuto {
		format = DetectFormat(data)
	}
	stats := ParseStats{Bytes: int64(len(data)), Format: format}
	var (
		rules []Rule
		err   error
	)
	switch format {
	case FormatHosts:
		rules, stats, err = parseHosts(data, opt)
	case FormatAdblock:
		rules, stats, err = parseAdblock(data, opt)
	case FormatPlain:
		rules, stats, err = parsePlain(data, opt)
	case FormatRPZ:
		rules, stats, err = parseRPZ(data, opt)
	default:
		return nil, stats, fmt.Errorf("blocklist: unknown format %q", format)
	}
	return rules, stats, err
}

func parseHosts(data []byte, opt ParseOptions) ([]Rule, ParseStats, error) {
	var rules []Rule
	stats := ParseStats{Bytes: int64(len(data)), Format: FormatHosts}
	sc := newLineScanner(data)
	for sc.Scan() {
		line := sc.Text()
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			continue
		}
		if trimmed[0] == '#' || trimmed[0] == ';' {
			stats.Comments++
			continue
		}
		// Strip inline comments.
		if i := strings.IndexByte(trimmed, '#'); i >= 0 {
			trimmed = strings.TrimSpace(trimmed[:i])
			if trimmed == "" {
				stats.Comments++
				continue
			}
		}
		fields := strings.Fields(trimmed)
		if len(fields) < 2 {
			stats.Malformed++
			continue
		}
		addr := fields[0]
		if !isAddr(addr) && !strings.EqualFold(addr, "localhost") {
			// Not an address: this is not a hosts file after all, so treat
			// the whole line as a domain list. The notice has to travel with
			// the stats the plain parser returns, or the operator sees the
			// format switch with no explanation for it.
			const notice = "no address in first field, falling back to a plain domain list"
			rules, pstats, perr := parsePlain(data, opt)
			pstats.Notices = appendNotice(pstats.Notices, notice)
			return rules, pstats, perr
		}
		names := fields[1:]
		if isLoopbackName(names) {
			// "127.0.0.1 localhost" style bootstrap lines: never block these.
			stats.Skipped++
			continue
		}
		for _, n := range names {
			if !validDomain(n) {
				stats.Malformed++
				continue
			}
			rules = append(rules, hostsRule(n, opt))
		}
	}
	if err := sc.Err(); err != nil {
		return rules, stats, err
	}
	stats.Rules = len(rules)
	return rules, stats, nil
}

func isLoopbackName(names []string) bool {
	for _, n := range names {
		switch strings.ToLower(strings.TrimSuffix(n, ".")) {
		case "localhost", "localhost.localdomain", "broadcasthost", "ip6-localhost", "ip6-loopback":
			return true
		}
	}
	return false
}

// hostsRule builds a rule for a hosts-file or plain-list entry.
//
// A hosts file says "this name is bad" and everyone means the name *and*
// anything under it, so the pattern is "**." (apex included) rather than the
// strict-subdomain "*.". Getting this wrong is how a filter ends up blocking
// ads.example.com but happily resolving 1.ads.example.com.
func hostsRule(name string, opt ParseOptions) Rule {
	pattern := name
	if opt.Subdomains {
		pattern = "**." + name
	}
	return Rule{Name: pattern, Action: ActionBlock, TTL: opt.DefaultTTL}
}

func parseAdblock(data []byte, opt ParseOptions) ([]Rule, ParseStats, error) {
	var rules []Rule
	stats := ParseStats{Bytes: int64(len(data)), Format: FormatAdblock}
	sc := newLineScanner(data)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" {
			continue
		}
		if line[0] == '!' || line[0] == '#' {
			stats.Comments++
			continue
		}
		if line[0] == '[' {
			stats.Comments++
			continue
		}
		if strings.HasPrefix(line, "||") {
			host := hostOfAdblock(line[2:])
			if !validDomain(host) {
				stats.Skipped++
				continue
			}
			rules = append(rules, Rule{Name: "**." + host, Action: ActionBlock, TTL: opt.DefaultTTL})
			continue
		}
		if strings.HasPrefix(line, "@@") {
			rest := strings.TrimPrefix(line, "@@")
			host := hostOfAdblock(strings.TrimPrefix(rest, "||"))
			if !validDomain(host) {
				stats.Skipped++
				continue
			}
			rules = append(rules, Rule{Name: "**." + host, Action: ActionAllow, Priority: 10, TTL: opt.DefaultTTL})
			continue
		}
		if strings.HasPrefix(line, "|") {
			rest := strings.TrimPrefix(line, "|")
			if strings.HasPrefix(rest, "http://") || strings.HasPrefix(rest, "https://") {
				host := hostOfURL(rest)
				if !validDomain(host) {
					stats.Skipped++
					continue
				}
				// A path rule cannot be expressed in DNS; blocking the host
				// is the closest faithful approximation, so record it once.
				rules = append(rules, Rule{Name: "**." + host, Action: ActionBlock, TTL: opt.DefaultTTL})
				continue
			}
			stats.Skipped++
			continue
		}
		if line[0] == '/' && strings.HasSuffix(line, "/") {
			stats.Skipped++ // regular expressions are not DNS-expressible
			if len(stats.Notices) < 5 {
				stats.Notices = appendNotice(stats.Notices, "skipped regular-expression rules")
			}
			continue
		}
		// Element-hiding rules ("##", "#@#", "#?#", "##+js(") carry no DNS
		// meaning at all.
		if strings.Contains(line, "##") || strings.Contains(line, "#@#") || strings.Contains(line, "#?#") {
			stats.Skipped++
			continue
		}
		// A bare hostname with no adblock syntax at all.
		if validDomain(line) {
			rules = append(rules, hostsRule(line, opt))
			continue
		}
		stats.Skipped++
	}
	if err := sc.Err(); err != nil {
		return rules, stats, err
	}
	stats.Rules = len(rules)
	return rules, stats, nil
}

// hostOfAdblock extracts the domain from the part of an adblock rule between
// "||" (or a scheme) and the first separator (^, /, ?, |, $).
func hostOfAdblock(s string) string {
	if i := strings.IndexAny(s, "^/?|$"); i >= 0 {
		s = s[:i]
	}
	return strings.Trim(s, ".")
}

func hostOfURL(s string) string {
	if i := strings.Index(s, "://"); i >= 0 {
		s = s[i+3:]
	}
	if i := strings.IndexAny(s, "/?#"); i >= 0 {
		s = s[:i]
	}
	if i := strings.LastIndexByte(s, '@'); i >= 0 {
		s = s[i+1:]
	}
	// Strip credentials and port.
	if i := strings.LastIndexByte(s, ':'); i >= 0 {
		if _, err := strconv.Atoi(s[i+1:]); err == nil {
			s = s[:i]
		}
	}
	return strings.Trim(s, ".")
}

func parsePlain(data []byte, opt ParseOptions) ([]Rule, ParseStats, error) {
	var rules []Rule
	stats := ParseStats{Bytes: int64(len(data)), Format: FormatPlain}
	sc := newLineScanner(data)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" {
			continue
		}
		if line[0] == '#' || line[0] == ';' {
			stats.Comments++
			continue
		}
		// Accept "0.0.0.0 domain" and bare "domain" alike so operators do
		// not have to normalise their sources by hand.
		if fields := strings.Fields(line); len(fields) > 1 && isAddr(fields[0]) {
			for _, n := range fields[1:] {
				if validDomain(n) {
					rules = append(rules, hostsRule(n, opt))
				} else {
					stats.Malformed++
				}
			}
			continue
		}
		// Support inline allow markers: the documented "!allow example.com"
		// and the bare "!example.com", plus the adblock exception form
		// "@@||example.com^" operators paste into plain lists too.
		allow := false
		if strings.HasPrefix(line, "@@") {
			allow = true
			line = strings.TrimPrefix(line, "@@")
		} else if strings.HasPrefix(line, "!") {
			allow = true
			line = trimAllowWord(strings.TrimSpace(line[1:]))
		}
		// A plain list often carries a few adblock lines; drop the marker and
		// everything from the first adblock separator on, so they survive
		// instead of being counted as malformed.
		line = cutAdblockTail(strings.TrimPrefix(line, "||"))
		if !validDomain(line) {
			stats.Malformed++
			continue
		}
		pattern := line
		if opt.Subdomains && !strings.HasPrefix(line, "*") {
			pattern = "**." + line
		}
		action := ActionBlock
		if allow {
			action = ActionAllow
		}
		rules = append(rules, Rule{Name: pattern, Action: action, TTL: opt.DefaultTTL})
	}
	if err := sc.Err(); err != nil {
		return rules, stats, err
	}
	stats.Rules = len(rules)
	return rules, stats, nil
}

// cutAdblockTail drops everything from the first adblock separator onwards, so
// a plain list that someone pasted adblock lines into still yields the domain
// instead of a name with "^$|*//" decoration glued to it. A leading "*." is a
// wildcard prefix rather than a separator, and a lone "$" separator is only
// decoration when something follows it.
func cutAdblockTail(s string) string {
	for i := 0; i < len(s); i++ {
		switch s[i] {
		case '^', '|', '/':
			return s[:i]
		case '$':
			if i+1 < len(s) {
				return s[:i]
			}
		case '*':
			if i > 0 {
				return s[:i]
			}
		}
	}
	return s
}

// trimAllowWord drops a leading allow keyword that is separated from the
// domain by whitespace, so the documented "!allow example.com" and the bare
// "!example.com" mean the same thing, while a domain that merely starts with
// the word (allow.example.com) is left alone.
func trimAllowWord(s string) string {
	fields := strings.Fields(s)
	if len(fields) < 2 {
		return s
	}
	switch strings.ToLower(fields[0]) {
	case "allow", "allowlist", "whitelist":
		return strings.TrimSpace(s[len(fields[0]):])
	}
	return s
}

func parseRPZ(data []byte, opt ParseOptions) ([]Rule, ParseStats, error) {
	var rules []Rule
	stats := ParseStats{Bytes: int64(len(data)), Format: FormatRPZ}
	z, err := zonefile.Parse(bytes.NewReader(data), zonefile.Options{DefaultTTL: opt.DefaultTTL})
	if err != nil {
		return nil, stats, err
	}
	// A record line with no TTL column comes back from the zone reader
	// carrying its own placeholder, which would shadow the TTL the caller
	// asked for, so the zone's own default wins for those records.
	defaultTTL := z.DefaultTTL
	if defaultTTL == 0 {
		defaultTTL = opt.DefaultTTL
	}
	seen := make(map[string]bool, len(z.Records))
	for _, rec := range z.Records {
		ttl := rec.TTL
		if ttl == 0 || (ttl == implicitRRTTL && defaultTTL != implicitRRTTL) {
			ttl = defaultTTL
		}
		rule, ok := rpzRule(rec, ttl, opt)
		if !ok {
			stats.Skipped++
			continue
		}
		if seen[rule.Name] {
			stats.Duplicates++
			continue
		}
		seen[rule.Name] = true
		rules = append(rules, rule)
	}
	stats.Rules = len(rules)
	return rules, stats, nil
}

// implicitRRTTL is the TTL the zone reader stamps on a record whose line
// carries no TTL column. It is a placeholder rather than the file's intent,
// which is why parseRPZ replaces it with the zone's own default.
const implicitRRTTL = 3600

// rpzRule maps one policy-zone record onto a rule (RFC draft-ietf-dnsop-rpz).
// RPZ encodes the response in the rdata: a CNAME to "." means NXDOMAIN, to
// "rpz-drop." means drop, to "rpz-passthru." means resolve normally, and
// A/AAAA records return that address verbatim.
func rpzRule(rec zonefile.Record, ttl uint32, opt ParseOptions) (Rule, bool) {
	pattern := rec.Name
	target, action, resp := rpzTarget(rec.RR)
	if target == "" {
		return Rule{}, false
	}
	if !strings.HasPrefix(pattern, "*") && opt.Subdomains {
		pattern = "*." + pattern
	}
	return Rule{Name: pattern, Action: action, Response: resp, TTL: ttl}, true
}

// rpzTarget interprets a policy record's rdata. The returned target is only a
// "is this actionable" flag; action and resp carry the meaning.
func rpzTarget(rr dns.RR) (target string, action Action, resp string) {
	switch v := rr.(type) {
	case *dns.CNAME:
		t := v.Target
		switch {
		case t == ".":
			// The classic RPZ "this name does not exist" trigger.
			return ".", ActionNXDomain, ""
		case strings.EqualFold(t, "rpz-drop."):
			return t, ActionDrop, ""
		case strings.EqualFold(t, "rpz-passthru."), strings.EqualFold(t, "rpz-passthru."):
			return t, ActionPassthrough, ""
		case strings.EqualFold(t, "rpz-tcp-only."):
			// Force TCP upstream; as a DNS policy this is a no-op for us
			// since we already retry over TCP on truncation.
			return t, ActionPassthrough, ""
		default:
			return t, ActionCNAME, t
		}
	case *dns.A:
		if v.A == nil {
			return "", ActionBlock, ""
		}
		return v.A.String(), ActionRedirect, v.A.String()
	case *dns.AAAA:
		if v.AAAA == nil {
			return "", ActionBlock, ""
		}
		return v.AAAA.String(), ActionRedirect, v.AAAA.String()
	case *dns.TXT:
		// Some providers publish policy in TXT records; the payload is the
		// action name.
		t := strings.ToLower(strings.Join(v.Txt, ""))
		switch {
		case strings.Contains(t, "rpz-drop"):
			return t, ActionDrop, ""
		case strings.Contains(t, "rpz-passthru"):
			return t, ActionPassthrough, ""
		case strings.Contains(t, "allow"), strings.Contains(t, "whitelist"):
			return t, ActionAllow, ""
		default:
			return "", ActionBlock, ""
		}
	default:
		return "", ActionBlock, ""
	}
}

// isAddr reports whether s parses as an IP address literal.
func isAddr(s string) bool {
	_, err := netip.ParseAddr(s)
	return err == nil
}

// newLineScanner returns a line scanner sized for the very large list files
// real blocklists ship (multi-hundred-megabyte entries are not unheard of).
func newLineScanner(data []byte) *bufio.Scanner {
	sc := bufio.NewScanner(bytes.NewReader(data))
	sc.Buffer(make([]byte, 0, 128*1024), 4*1024*1024)
	return sc
}

func appendNotice(notices []string, n string) []string {
	for _, existing := range notices {
		if existing == n {
			return notices
		}
	}
	if len(notices) >= 8 {
		return notices
	}
	return append(notices, n)
}

// validDomain is a cheap sanity check: no scheme, no path, printable, and at
// least one dot or a known single-label TLD-ish name.
func validDomain(s string) bool {
	if s == "" || len(s) > 253 {
		return false
	}
	if strings.HasPrefix(s, ".") || strings.HasPrefix(s, "-") {
		return false
	}
	if strings.ContainsAny(s, "/ \t\\\"'<>|?*=") {
		return false
	}
	labels := strings.Split(strings.TrimSuffix(s, "."), ".")
	if len(labels) < 1 {
		return false
	}
	for _, l := range labels {
		if l == "" || len(l) > 63 {
			return false
		}
		for i := 0; i < len(l); i++ {
			c := l[i]
			ok := c == '-' || c == '_' || c == '*' ||
				(c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
				(c >= '0' && c <= '9')
			if !ok {
				return false
			}
		}
		if l[0] == '-' || l[len(l)-1] == '-' {
			return false
		}
	}
	// A bare label with no dot is only meaningful for the localhost-ish
	// names, which callers already filter; reject to avoid blocking a
	// single-label intranet name by accident.
	return strings.Contains(s, ".")
}