"""Mapline — victim / actor / jurisdiction resolution.

Workers: GeoResolver + CVEResolver. GeoResolver resolves a case's primary host to a country code
via ip-api.com (free, no auth, 45 req/min). For malicious-infrastructure
cases (URLhaus etc.) "victim.country" carries the hosting-country semantic;
documented in psyc.lines.route's destination policy.
"""

from __future__ import annotations

import socket
from typing import Iterable, Optional, Set
from urllib.parse import urlparse

import httpx

from psyc import log
from psyc.models import Case, Severity
from psyc.result import Err, Ok, Result


_log = log.get(__name__)

IP_API_URL = "http://ip-api.com/json/{ip}?fields=status,country,countryCode,isp,as,message"
USER_AGENT = "psyc/0.1 (defensive CTI; hackathon prototype)"
DNS_TIMEOUT_SECONDS = 3.0


def resolve(case: Case) -> Case:
    if case.victim.country:
        return case
    ip = _case_ip(case)
    if not ip:
        _log.info("map.geo.no_target", case_id=case.case_id)
        return case
    country_result = _geoip_country(ip)
    if isinstance(country_result, Err):
        _log.info("map.geo.skip", case_id=case.case_id, ip=ip, reason=country_result.reason)
        return case
    case.victim.country = country_result.value
    if ip not in case.observables.ips:
        case.observables.ips.append(ip)
    _log.info("map.geo.resolved", case_id=case.case_id, ip=ip, country=country_result.value)
    return case


def _case_ip(case: Case) -> Optional[str]:
    if case.observables.ips:
        return case.observables.ips[0]
    if case.observables.domains:
        return _resolve_dns(case.observables.domains[0])
    if case.observables.urls:
        host = urlparse(case.observables.urls[0]).hostname
        if host:
            return _resolve_dns(host)
    return None


def _resolve_dns(host: str) -> Optional[str]:
    if _looks_like_ip(host):
        return host
    prev_timeout = socket.getdefaulttimeout()
    socket.setdefaulttimeout(DNS_TIMEOUT_SECONDS)
    try:
        return socket.gethostbyname(host)
    except (socket.gaierror, socket.herror, socket.timeout):
        return None
    finally:
        socket.setdefaulttimeout(prev_timeout)


def _looks_like_ip(host: str) -> bool:
    parts = host.split(".")
    return len(parts) == 4 and all(p.isdigit() and 0 <= int(p) < 256 for p in parts)


def _geoip_country(ip: str, timeout: float = 5.0) -> Result[str, str]:
    try:
        with httpx.Client(timeout=timeout, headers={"User-Agent": USER_AGENT}) as client:
            r = client.get(IP_API_URL.format(ip=ip))
            r.raise_for_status()
            body = r.json()
    except httpx.HTTPError as exc:
        return Err(f"ip-api request failed: {exc}")
    if body.get("status") != "success":
        return Err(f"ip-api: {body.get('message', 'unknown error')}")
    code = body.get("countryCode")
    if not code:
        return Err("ip-api returned no countryCode")
    return Ok(code)


# --- CVEResolver — cross-check case CVEs against the CISA KEV catalog --------

def kev_cve_set(cases: Iterable[Case]) -> Set[str]:
    """The set of CVE IDs carried by CISA KEV cases — the known-exploited catalog."""
    out: Set[str] = set()
    for case in cases:
        if case.source_metadata.get("feed") == "cisa-kev":
            out.update(cve.upper() for cve in case.observables.cves)
    return out


def resolve_cves(case: Case, kev_cves: Set[str]) -> Case:
    """Flag any of the case's CVEs that are known-exploited; escalate if so."""
    if not case.observables.cves:
        return case
    hits = sorted(c for c in case.observables.cves if c.upper() in kev_cves)
    if not hits:
        return case
    case.source_metadata["kev_cves"] = ",".join(hits)
    # a known-exploited CVE surfacing on a non-KEV case is a real escalation
    if case.source_metadata.get("feed") != "cisa-kev":
        if case.classification.severity in (None, Severity.LOW, Severity.MEDIUM):
            case.classification.severity = Severity.HIGH
    _log.info("map.cve.kev_match", case_id=case.case_id, kev_cves=hits)
    return case
