"""Persist results and artifacts produced by the shared analysis context."""

from __future__ import annotations

import json
import logging
import os

from django.conf import settings
from analyzer import AnalysisEngine, extract_urls, extract_hashes
from analyzer.extract_pyshark import build_context_from_pcap
from analyzer.findings import Severity
from IoTAnalyzer.models import FindingReport, UrlRaw, HashRaw

logger = logging.getLogger(__name__)

_CANDIDATE_DIRS = [
    str(settings.CAPTURE_UPLOAD_DIR),
]


def resolve_capture_path(filename: str) -> str | None:
    """Resolve a stored capture name against supported upload locations."""
    if not filename:
        return None
    if os.path.isabs(filename) and os.path.exists(filename):
        return filename
    for base in _CANDIDATE_DIRS:
        candidate = os.path.join(base, filename)
        if os.path.exists(candidate):
            return candidate
    return None


def _max_severity(result) -> str:
    worst = Severity.INFO
    for f in result.findings:
        if f.severity.rank > worst.rank:
            worst = f.severity
    return worst.value


def collect_cloud_results(file_id: int) -> dict:
    """Return deduplicated capture artifacts and any enrichment results."""
    from IoTAnalyzer.models import UrlRaw, HashRaw, UrlAnalyzed, HashAnalyzed

    # Retain raw IDs so the frontend can request a selected enrichment subset.
    raw_urls = list(UrlRaw.objects.filter(raw_dataset__file_id=file_id)
                    .values_list('id', 'url'))
    raw_hashes = list(HashRaw.objects.filter(raw_dataset__file_id=file_id)
                      .values_list('id', 'hash_value'))

    url_map = {}
    for ua in (UrlAnalyzed.objects.filter(analyzed_dataset__file_id=file_id)
               .prefetch_related('ports', 'tls_versions', 'cipher_suites')):
        url_map[ua.url] = {
            "ports": sorted({p.port for p in ua.ports.all()}),
            "tls_versions": [t.tls_version for t in ua.tls_versions.all()],
            "cipher_suites": [c.cipher_suite for c in ua.cipher_suites.all()],
        }

    hash_map = {}
    for ha in (HashAnalyzed.objects.filter(analyzed_dataset__file_id=file_id)
               .prefetch_related('hash_type')):
        entry = {
            "cracked": ha.cracked,
            "cracked_value": ha.cracked_value,
            "crack_time": ha.crack_time,
            "hash_type": [h.hash_type for h in ha.hash_type.all()],
        }
        existing = hash_map.get(ha.hash_value)
        if existing is None or (ha.cracked and not existing.get("cracked")):
            hash_map[ha.hash_value] = entry

    urls, seen = [], set()
    for raw_id, u in raw_urls:
        if not u or u in seen:
            continue
        seen.add(u)
        a = url_map.get(u)
        urls.append({"id": raw_id, "url": u, "analyzed": a is not None, **(a or {})})

    hashes, seen_h = [], set()
    for raw_id, h in raw_hashes:
        if not h or h in seen_h:
            continue
        seen_h.add(h)
        a = hash_map.get(h)
        hashes.append({"id": raw_id, "hash_value": h, "analyzed": a is not None,
                       **(a or {"cracked": False, "cracked_value": None,
                                "crack_time": None, "hash_type": []})})

    return {"urls": urls, "hashes": hashes}


def _persist_artifacts(ctx, raw_dataset_id: int) -> None:
    """Persist URL and hash artifacts extracted during the analysis pass."""
    if raw_dataset_id is None:
        return
    try:
        for u in extract_urls(ctx):
            UrlRaw.objects.get_or_create(raw_dataset_id=raw_dataset_id, url=u.value[:200])
    except Exception as exc:
        logger.exception("URL persistence failed: %s", exc)
    try:
        for h in extract_hashes(ctx):
            HashRaw.objects.get_or_create(raw_dataset_id=raw_dataset_id, hash_value=h.value)
    except Exception as exc:
        logger.exception("Hash persistence failed: %s", exc)


def analyze_capture(filename: str, file_id: int, raw_dataset_id: int | None = None):
    """Analyze a capture once and persist findings plus extracted artifacts."""
    path = resolve_capture_path(filename)
    if not path:
        logger.warning("Capture file not found: %s", filename)
        return None

    try:
        ctx = build_context_from_pcap(path)
        result = AnalysisEngine().analyze_context(ctx)
    except Exception as exc:
        logger.exception("Security analysis failed for %s: %s", filename, exc)
        return None

    try:
        # Reanalysis replaces the previous report atomically at the file level.
        FindingReport.objects.filter(file_id=file_id).delete()
        FindingReport.objects.create(
            file_id=file_id,
            result_json=json.dumps(result.to_dict()),
            finding_count=len(result.findings),
            max_severity=_max_severity(result),
        )
        logger.info("Stored finding report: file_id=%s, findings=%s",
                    file_id, len(result.findings))
    except Exception as exc:
        logger.exception("Finding report persistence failed: %s", exc)

    _persist_artifacts(ctx, raw_dataset_id)

    return result


def run_security_analysis(filename: str, file_id: int):
    return analyze_capture(filename, file_id, raw_dataset_id=None)