from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, List, Optional

from bastion.models import DriftEvent, ScanResult

ESCALATION_CATEGORIES = {"system_execution", "dynamic_code", "obfuscation"}


def load_baseline(bastion_dir: Path) -> Optional[List[Dict[str, Any]]]:
    path = bastion_dir / "baseline.json"
    if not path.exists():
        return None
    with open(path, "r") as f:
        data = json.load(f)
    return data.get("results", [])


def save_baseline(scan_results: List[ScanResult], bastion_dir: Path, label: str = "Baseline") -> None:
    path = bastion_dir / "baseline.json"
    data = {
        "label": label,
        "results": [r.to_dict() for r in scan_results],
    }
    with open(path, "w") as f:
        json.dump(data, f, indent=2)


def _has_escalation_capability(capabilities: List[Dict[str, Any]]) -> bool:
    for cap in capabilities:
        if cap.get("category") in ESCALATION_CATEGORIES:
            return True
    return False


def _escalate_severity(base_severity: str, capabilities: List[Dict[str, Any]]) -> str:
    if _has_escalation_capability(capabilities):
        cap_severities = [cap.get("severity", "Low") for cap in capabilities]
        if "Critical" in cap_severities:
            return "Critical"
        return "High"
    return base_severity


def compute_drift(current_results: List[ScanResult], baseline: List[Dict[str, Any]]) -> List[DriftEvent]:
    baseline_map: Dict[str, Dict[str, Any]] = {r["plugin_id"]: r for r in baseline}
    current_map: Dict[str, ScanResult] = {r.plugin_id: r for r in current_results}

    drift_events: List[DriftEvent] = []

    for plugin_id, result in current_map.items():
        if plugin_id not in baseline_map:
            caps = [c.to_dict() for c in result.capabilities]
            severity = _escalate_severity("Medium", caps)
            drift_events.append(DriftEvent(
                plugin_id=plugin_id,
                event_type="new",
                capability_changes=[{"added": c} for c in caps],
                severity=severity,
            ))
        else:
            old_caps = set(c["category"] for c in baseline_map[plugin_id].get("capabilities", []))
            new_caps = set(c.category for c in result.capabilities)
            added = new_caps - old_caps
            removed = old_caps - new_caps

            if added or removed:
                changes: List[Dict[str, Any]] = []
                for cat in added:
                    changes.append({"type": "added", "category": cat})
                for cat in removed:
                    changes.append({"type": "removed", "category": cat})

                all_caps = [c.to_dict() for c in result.capabilities]
                severity = _escalate_severity("Medium", all_caps)
                drift_events.append(DriftEvent(
                    plugin_id=plugin_id,
                    event_type="modified",
                    capability_changes=changes,
                    severity=severity,
                ))

    for plugin_id in baseline_map:
        if plugin_id not in current_map:
            drift_events.append(DriftEvent(
                plugin_id=plugin_id,
                event_type="removed",
                capability_changes=[],
                severity="Low",
            ))

    return drift_events
