"""Classifyline — assigns incident_type, severity, TLP, internal_class."""

from __future__ import annotations

from psyc import log
from psyc.models import Case, IncidentType, InternalClass, Severity, TLP


_log = log.get(__name__)

_FEED_INCIDENT = {
    "urlhaus": IncidentType.MALWARE,
    "feodo": IncidentType.BOTNET,
    "cisa-kev": IncidentType.EXPLOIT,
    "malware-bazaar": IncidentType.MALWARE,
    "otx": IncidentType.MALWARE,  # default; OTX pulses span many types
}

# ThreatFox carries its own type signal — map it instead of using a feed default.
_THREATFOX_THREAT_TYPE = {
    "botnet_cc": IncidentType.BOTNET,
    "payload_delivery": IncidentType.MALWARE,
    "payload": IncidentType.MALWARE,
    "phishing": IncidentType.PHISHING,
}


def classify(case: Case) -> Case:
    _classify_incident_type_and_tlp(case)
    _classify_severity(case)
    _classify_internal_class(case)
    _log.info(
        "classify.applied",
        case_id=case.case_id,
        incident_type=case.classification.incident_type.value if case.classification.incident_type else None,
        severity=case.classification.severity.value if case.classification.severity else None,
        tlp=case.classification.tlp.value,
        internal_class=case.classification.internal_class.value if case.classification.internal_class else None,
    )
    return case


def _classify_incident_type_and_tlp(case: Case) -> None:
    if case.classification.incident_type is not None:
        return
    feed = case.source_metadata.get("feed", "")
    if feed == "threatfox":
        incident = _THREATFOX_THREAT_TYPE.get(case.source_metadata.get("threat_type", ""), IncidentType.MALWARE)
    else:
        incident = _FEED_INCIDENT.get(feed)
    if incident is None and case.observables.urls:
        incident = IncidentType.MALWARE  # fallback for un-tagged feeds
    if incident is None:
        return
    case.classification.incident_type = incident
    if case.classification.tlp == TLP.AMBER:
        case.classification.tlp = TLP.GREEN


def _classify_severity(case: Case) -> None:
    if case.classification.severity is not None:
        return
    if case.victim.critical_infrastructure:
        case.classification.severity = Severity.CRITICAL
        return
    incident = case.classification.incident_type
    if incident == IncidentType.EXPLOIT:
        ransomware = case.source_metadata.get("ransomware", "")
        case.classification.severity = Severity.CRITICAL if ransomware == "Known" else Severity.HIGH
        return
    if incident in (IncidentType.MALWARE, IncidentType.BOTNET):
        status = case.source_metadata.get("url_status") or case.source_metadata.get("status", "")
        case.classification.severity = Severity.HIGH if status == "online" else Severity.MEDIUM
        return
    case.classification.severity = Severity.MEDIUM


def _classify_internal_class(case: Case) -> None:
    if case.classification.internal_class is not None:
        return
    if case.victim.critical_infrastructure or case.classification.severity == Severity.CRITICAL:
        case.classification.internal_class = InternalClass.A
        return
    if case.classification.severity == Severity.HIGH:
        case.classification.internal_class = InternalClass.C
        return
    case.classification.internal_class = InternalClass.D
