import Foundation

/// A grouping of destinations under one real-world owner. `key` is the stable
/// grouping id (a registrable domain, or a bare IP when unresolved); `name` is
/// what the district is labelled.
public struct Org: Sendable, Hashable {
    public let key: String
    public let name: String
    public let resolvedHost: String?

    public init(key: String, name: String, resolvedHost: String?) {
        self.key = key
        self.name = name
        self.resolvedHost = resolvedHost
    }
}

/// Maps a reverse-DNS hostname to an `Org`. Pure and deterministic so it can be
/// unit-tested without touching the network.
public enum OrgClassifier {

    private struct Rule { let key: String; let name: String; let domains: Set<String> }

    /// Well-known owners whose traffic spreads across many registrable domains.
    /// Anything matched here collapses to a single district.
    private static let rules: [Rule] = [
        Rule(key: "apple", name: "Apple",
             domains: ["apple.com", "icloud.com", "aaplimg.com", "mzstatic.com", "cdn-apple.com", "apple-dns.net"]),
        Rule(key: "google", name: "Google",
             domains: ["google.com", "googleusercontent.com", "gstatic.com", "googleapis.com",
                       "1e100.net", "ggpht.com", "youtube.com", "ytimg.com", "google-analytics.com"]),
        Rule(key: "amazon", name: "Amazon / AWS",
             domains: ["amazonaws.com", "amazon.com", "cloudfront.net", "aws.dev", "awsglobalaccelerator.com"]),
        Rule(key: "microsoft", name: "Microsoft",
             domains: ["microsoft.com", "windows.com", "windowsupdate.com", "azure.com",
                       "azureedge.net", "office.com", "live.com", "msftncsi.com"]),
        Rule(key: "meta", name: "Meta",
             domains: ["facebook.com", "fbcdn.net", "instagram.com", "whatsapp.net", "fb.com"]),
        Rule(key: "cloudflare", name: "Cloudflare",
             domains: ["cloudflare.com", "cloudflare-dns.com", "cf-dns.com"]),
        Rule(key: "akamai", name: "Akamai",
             domains: ["akamai.net", "akamaiedge.net", "akamaitechnologies.com", "akadns.net"]),
        Rule(key: "fastly", name: "Fastly",
             domains: ["fastly.net", "fastlylb.net"]),
        Rule(key: "github", name: "GitHub",
             domains: ["github.com", "githubusercontent.com", "githubassets.com"]),
        Rule(key: "spotify", name: "Spotify",
             domains: ["spotify.com", "scdn.co", "spotifycdn.com"]),
        Rule(key: "netflix", name: "Netflix",
             domains: ["netflix.com", "nflxvideo.net", "nflxso.net", "nflximg.net"]),
        Rule(key: "telegram", name: "Telegram",
             domains: ["telegram.org", "t.me", "telegram.me"]),
        Rule(key: "anthropic", name: "Anthropic",
             domains: ["anthropic.com", "claude.ai"]),
    ]

    /// Registrable suffixes that span two labels, so the registrable domain is
    /// the last *three* labels (e.g. `bbc.co.uk`, not `co.uk`).
    private static let multiPartSuffixes: Set<String> = [
        "co.uk", "org.uk", "gov.uk", "ac.uk", "co.jp", "co.nz", "co.in",
        "com.au", "net.au", "org.au", "com.br", "com.cn", "com.mx", "co.kr", "co.za",
    ]

    public static func classify(ip: String, ptr: String?) -> Org {
        // 1. A PTR with a real registrable domain is the most accurate signal.
        if let ptr, let reg = registrableDomain(ptr) {
            if let rule = rules.first(where: { $0.domains.contains(reg) }) {
                return Org(key: rule.key, name: rule.name, resolvedHost: ptr)
            }
            return Org(key: reg, name: prettify(reg), resolvedHost: ptr)
        }
        // 2. No usable PTR — fall back to the curated CIDR table (Apple's 17/8,
        //    Telegram, private LAN, …).
        if let match = IPRanges.match(ip) {
            return Org(key: match.key, name: match.name, resolvedHost: ptr)
        }
        // 3. Unknown: the IP is its own (unlabelled) district.
        return Org(key: ip, name: ip, resolvedHost: ptr)
    }

    /// eTLD+1 with a small built-in multi-part-suffix table. `a.b.example.co.uk`
    /// -> `example.co.uk`; `cdn.gstatic.com` -> `gstatic.com`.
    public static func registrableDomain(_ host: String) -> String? {
        let trimmed = host.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: "."))
        let labels = trimmed.split(separator: ".").map(String.init)
        guard labels.count >= 2 else { return nil }
        let lastTwo = labels.suffix(2).joined(separator: ".")
        if labels.count >= 3 && multiPartSuffixes.contains(lastTwo) {
            return labels.suffix(3).joined(separator: ".")
        }
        return lastTwo
    }

    /// `example.co.uk` -> `Example`.
    public static func prettify(_ registrable: String) -> String {
        guard let sld = registrable.split(separator: ".").first else { return registrable }
        return sld.prefix(1).uppercased() + sld.dropFirst()
    }
}