package privacy

import (
	"sort"
	"strconv"
	"strings"

	"git.mchus.pro/mchus/logpile/internal/models"
)

// Multi-label public suffixes handled explicitly; every other final label is
// treated as the suffix. This is a heuristic, not a full PSL - good enough to
// collapse "sn-x.mgmt.corp.example.co.uk" to "example.co.uk".
var multiLabelSuffixes = []string{
	"co.uk", "org.uk", "gov.uk", "ac.uk", "me.uk",
	"com.cn", "net.cn", "org.cn", "gov.cn", "edu.cn",
	"com.au", "net.au", "org.au", "co.jp", "co.kr", "com.br", "com.tr", "com.ua",
}

// Categories whose findings feed the customer guess, and whether a hit in that
// category is strong evidence on its own.
var customerCategories = map[string]bool{
	catDomain:        false,
	catResolv:        true,
	catADLDAP:        true,
	catCert:          true,
	catNSUpdate:      true,
	catMgmtSubdomain: true,
	catDHCP:          false,
	catEmail:         false,
}

type customerTally struct {
	domain       string
	hits         int
	paths        map[string]struct{}
	strong       bool
	evidence     []string
	evidenceSeen map[string]struct{}
}

func guessCustomers(findings []models.PrivacyFinding) []models.CustomerGuess {
	tallies := map[string]*customerTally{}

	for _, f := range findings {
		strongCat, ok := customerCategories[f.Category]
		if !ok {
			continue
		}
		host := f.Match
		if f.Category == catEmail {
			if at := strings.IndexByte(host, '@'); at >= 0 {
				host = host[at+1:]
			}
		}
		reg := registrableDomain(host)
		if reg == "" || isAllowlistedDomain(reg) {
			continue
		}
		t := tallies[reg]
		if t == nil {
			t = &customerTally{domain: reg, paths: map[string]struct{}{}, evidenceSeen: map[string]struct{}{}}
			tallies[reg] = t
		}
		t.hits++
		if f.Path != "" {
			t.paths[f.Path] = struct{}{}
		}
		if strongCat {
			t.strong = true
		}
		if len(t.evidence) < 3 && f.Excerpt != "" {
			line := evidenceLine(f)
			if _, dup := t.evidenceSeen[line]; !dup {
				t.evidenceSeen[line] = struct{}{}
				t.evidence = append(t.evidence, line)
			}
		}
	}

	guesses := make([]models.CustomerGuess, 0, len(tallies))
	for _, t := range tallies {
		guesses = append(guesses, models.CustomerGuess{
			Domain:     t.domain,
			Confidence: confidence(t),
			Hits:       t.hits,
			Evidence:   t.evidence,
		})
	}

	sort.Slice(guesses, func(i, j int) bool {
		si, sj := score(tallies[guesses[i].Domain]), score(tallies[guesses[j].Domain])
		if si != sj {
			return si > sj
		}
		return guesses[i].Domain < guesses[j].Domain
	})

	// Keep the strongest few; drop weak single-hit noise unless it is all we have.
	out := make([]models.CustomerGuess, 0, 3)
	for _, g := range guesses {
		if len(out) >= 3 {
			break
		}
		if g.Hits < 2 && g.Confidence == "low" && len(out) > 0 {
			continue
		}
		out = append(out, g)
	}
	return out
}

func score(t *customerTally) int {
	s := t.hits + 3*len(t.paths)
	if t.strong {
		s += 10
	}
	return s
}

func confidence(t *customerTally) string {
	switch {
	case t.strong && len(t.paths) >= 2:
		return "high"
	case t.strong || len(t.paths) >= 2:
		return "medium"
	default:
		return "low"
	}
}

func evidenceLine(f models.PrivacyFinding) string {
	loc := f.Path
	if f.Line > 0 {
		loc += ":" + strconv.Itoa(f.Line)
	}
	return loc + ": " + f.Excerpt
}

// registrableDomain returns the registrable ("eTLD+1") portion of a host name,
// or "" if host is an IP, has no dot, or is a bare TLD.
func registrableDomain(host string) string {
	host = strings.ToLower(strings.Trim(strings.TrimSpace(host), ".*"))
	host = strings.TrimPrefix(host, "*.")
	if host == "" || strings.ContainsAny(host, ":/ ") {
		return ""
	}
	if strings.Count(host, ".") == 3 {
		allDigits := true
		for _, r := range host {
			if r != '.' && (r < '0' || r > '9') {
				allDigits = false
				break
			}
		}
		if allDigits {
			return ""
		}
	}
	labels := strings.Split(host, ".")
	if len(labels) < 2 {
		return ""
	}
	for _, suf := range multiLabelSuffixes {
		if strings.HasSuffix(host, "."+suf) {
			sufLabels := strings.Count(suf, ".") + 1
			if len(labels) > sufLabels {
				return strings.Join(labels[len(labels)-sufLabels-1:], ".")
			}
			return ""
		}
	}
	return strings.Join(labels[len(labels)-2:], ".")
}