from collections import Counter
from ipaddress import ip_address

# ==============================
# Known Service Map (Expanded)
# ==============================
COMMON_PORTS = {
    20: "FTP-DATA",
    21: "FTP",
    22: "SSH",
    23: "TELNET",
    25: "SMTP",
    53: "DNS",
    67: "DHCP",
    68: "DHCP",
    69: "TFTP",
    80: "HTTP",
    110: "POP3",
    123: "NTP",
    137: "NetBIOS",
    138: "NetBIOS",
    139: "NetBIOS",
    143: "IMAP",
    161: "SNMP",
    389: "LDAP",
    443: "HTTPS",
    445: "SMB",
    465: "SMTPS",
    500: "ISAKMP",
    587: "SMTP-Submission",
    636: "LDAPS",
    989: "FTPS-DATA",
    990: "FTPS",
    1194: "OpenVPN",
    1433: "MSSQL",
    1521: "Oracle",
    2049: "NFS",
    2083: "cPanel",
    2087: "WHM",
    3000: "Dev-App",
    3306: "MySQL",
    3389: "RDP",
    4444: "Remote-Admin",
    5432: "PostgreSQL",
    5601: "Kibana",
    8000: "HTTP-Alt",
    8006: "Proxmox",
    8080: "HTTP-Proxy",
    8443: "HTTPS-Alt",
    9000: "Web-App",
}


# ==============================
# Helper Functions
# ==============================
def is_private(ip: str) -> bool:
    try:
        return ip_address(ip).is_private
    except Exception:
        return False


def classify_high_port(dport: int, external: bool):
    if dport >= 30000 and external:
        return {
            "service": "High-Port External Service",
            "guess": "Likely tunneling, VPN, or custom application",
            "confidence": 0.72,
        }

    if dport >= 30000 and not external:
        return {
            "service": "High-Port Internal Service",
            "guess": "Likely ephemeral or internal application",
            "confidence": 0.65,
        }

    return None


# ==============================
# Main Extractor
# ==============================
def extract_ports_and_services(packets):
    ports = Counter()
    services = Counter()
    connections = []

    for pkt in packets:
        if not pkt.haslayer("IP"):
            continue

        if not (pkt.haslayer("TCP") or pkt.haslayer("UDP")):
            continue

        layer = pkt["TCP"] if pkt.haslayer("TCP") else pkt["UDP"]

        src = pkt["IP"].src
        dst = pkt["IP"].dst
        sport = layer.sport
        dport = layer.dport

        ports[dport] += 1

        src_private = is_private(src)
        dst_private = is_private(dst)
        external = src_private and not dst_private

        # ---------- Known service ----------
        if dport in COMMON_PORTS:
            service = COMMON_PORTS[dport]
            guess = None
            confidence = 0.95

        else:
            # ---------- Heuristic detection ----------
            heuristic = classify_high_port(dport, external)
            if heuristic:
                service = heuristic["service"]
                guess = heuristic["guess"]
                confidence = heuristic["confidence"]
            else:
                service = "UNKNOWN"
                guess = "Insufficient data to classify"
                confidence = 0.30

        services[service] += 1

        connections.append({
            "src": src,
            "sport": sport,
            "dst": dst,
            "dport": dport,
            "service": service,
            "guess": guess,
            "confidence": confidence,
            "scope": "external" if external else "internal",
        })

    return {
        "ports": ports,
        "services": services,
        "connections": connections,
    }
