"""Thread-safe in-memory storage for security scan results and sessions."""

import hashlib
import json
import os
import re
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from maeris_mcp.types.security import (
    CheckResult,
    ScanSummary,
    SecurityFinding,
    SecurityScanResult,
    SeverityCounts,
    StoredSecurityScan,
)


@dataclass
class ScanMeta:
    """Additional metadata required for a scan session."""

    scan_root: str
    base_for_rel: str
    # Change-aware scanning: file content hashes from the previous scan and the current scan.
    # prev_file_hashes is loaded at scan creation; current_file_hashes is built lazily
    # during pagination and persisted at finalize time.
    prev_file_hashes: dict[str, str] = field(default_factory=dict)
    current_file_hashes: dict[str, str] = field(default_factory=dict)


@dataclass
class ScanIndexes:
    """Hashmap indexes for fast lookup of findings."""

    by_file: dict[str, list[str]] = field(default_factory=dict)
    by_rule: dict[str, list[str]] = field(default_factory=dict)
    by_owasp: dict[str, list[str]] = field(default_factory=dict)
    by_severity: dict[str, list[str]] = field(default_factory=dict)
    by_id: dict[str, SecurityFinding] = field(default_factory=dict)
    fingerprints: set[str] = field(default_factory=set)
    # File-level coverage tracking
    files_served: set[str] = field(default_factory=set)
    files_acknowledged: set[str] = field(default_factory=set)
    last_served_batch: list[str] = field(default_factory=list)


class SecurityScanStore:
    """Thread-safe storage for security scan results."""

    def __init__(self, max_scans: int = 100) -> None:
        self._lock = threading.RLock()
        self._scans: dict[str, StoredSecurityScan] = {}
        self._meta: dict[str, ScanMeta] = {}
        self._indexes: dict[str, ScanIndexes] = {}
        self._max_scans = max_scans

    def store(self, result: SecurityScanResult) -> str:
        """Store a scan result and return its ID."""
        return self.create_scan(
            result=result,
            file_list=[],
            selected_checks=[],
            scan_profile=None,
            scan_root=result.target_path,
            base_for_rel=result.target_path,
        )

    def create_scan(
        self,
        result: SecurityScanResult,
        file_list: list[str],
        selected_checks: list[str],
        scan_profile: str | list[str] | None,
        scan_root: str,
        base_for_rel: str,
    ) -> str:
        """Create a scan session and store its initial result."""
        with self._lock:
            self._enforce_max_size()
            now = datetime.now(timezone.utc).isoformat()
            stored = StoredSecurityScan(
                id=result.scan_id,
                target_path=result.target_path,
                result=result,
                stored_at=now,
                last_accessed=now,
                status="pending",
                selected_checks=selected_checks,
                scan_profile=scan_profile,
                file_list=file_list,
                pushed_to_maeris=False,
                pushed_at=None,
            )
            self._scans[result.scan_id] = stored
            self._meta[result.scan_id] = ScanMeta(scan_root=scan_root, base_for_rel=base_for_rel)
            self._indexes[result.scan_id] = ScanIndexes()
            return result.scan_id

    def get(self, scan_id: str) -> StoredSecurityScan | None:
        """Retrieve a scan by ID."""
        with self._lock:
            if scan := self._scans.get(scan_id):
                scan.last_accessed = datetime.now(timezone.utc).isoformat()
                return scan
            return None

    def get_all(self) -> list[StoredSecurityScan]:
        """Return all stored scans."""
        with self._lock:
            return list(self._scans.values())

    def get_meta(self, scan_id: str) -> ScanMeta | None:
        """Get additional scan metadata."""
        with self._lock:
            return self._meta.get(scan_id)

    def search_findings(
        self,
        scan_id: str,
        *,
        file_path: str | None = None,
        rule_id: str | None = None,
        owasp: str | None = None,
        severity: str | None = None,
    ) -> list[SecurityFinding]:
        """Search findings by indexed fields (internal use)."""
        with self._lock:
            indexes = self._indexes.get(scan_id)
            if not indexes:
                return []

            id_sets: list[set[str]] = []
            if file_path:
                id_sets.append(set(indexes.by_file.get(file_path, [])))
            if rule_id:
                id_sets.append(set(indexes.by_rule.get(rule_id, [])))
            if owasp:
                key = owasp.split(":", 1)[0]
                id_sets.append(set(indexes.by_owasp.get(key, [])))
            if severity:
                id_sets.append(set(indexes.by_severity.get(severity.lower(), [])))

            if not id_sets:
                return list(indexes.by_id.values())

            result_ids = set.intersection(*id_sets) if id_sets else set()
            return [indexes.by_id[i] for i in result_ids if i in indexes.by_id]

    def add_findings(self, scan_id: str, findings: list[SecurityFinding]) -> dict[str, int]:
        """Add findings to a scan, updating indexes and de-duplicating."""
        with self._lock:
            scan = self._scans.get(scan_id)
            indexes = self._indexes.get(scan_id)
            if not scan or not indexes:
                return {"added": 0, "deduped": 0}
            if scan.status == "completed":
                return {"added": 0, "deduped": 0}

            added = 0
            deduped = 0
            for finding in findings:
                # Drop findings with no file path — they can't be located,
                # validated, or acted upon, and break location-based deduplication.
                if not finding.file_path or not finding.file_path.strip():
                    deduped += 1
                    continue
                fingerprint = self._fingerprint(finding)
                if fingerprint in indexes.fingerprints:
                    deduped += 1
                    continue
                indexes.fingerprints.add(fingerprint)
                scan.result.findings.append(finding)
                self._index_finding(indexes, finding)
                added += 1
            return {"added": added, "deduped": deduped}

    def set_check_results(self, scan_id: str, results: list[CheckResult]) -> int:
        """Store compliance check results (pass/fail/not_applicable) for a scan."""
        with self._lock:
            scan = self._scans.get(scan_id)
            if not scan:
                return 0
            # Deduplicate by check_id — last result wins
            existing = {cr.check_id: cr for cr in scan.check_results}
            for cr in results:
                existing[cr.check_id] = cr
            scan.check_results = list(existing.values())
            return len(results)

    def record_files_served(self, scan_id: str, file_paths: list[str]) -> None:
        """Record that these file paths were returned in a fileContents batch."""
        with self._lock:
            if idx := self._indexes.get(scan_id):
                idx.last_served_batch = list(file_paths)
                idx.files_served.update(file_paths)

    def acknowledge_last_batch(self, scan_id: str) -> None:
        """Mark the last served batch as acknowledged (LLM submitted findings for it, even if empty)."""
        with self._lock:
            if idx := self._indexes.get(scan_id):
                idx.files_acknowledged.update(idx.last_served_batch)

    def mark_files_auto_processed(self, scan_id: str, file_paths: list[str]) -> None:
        """Mark files as both served and acknowledged without touching last_served_batch.

        Use for files the server processes automatically (e.g. dep-scan manifest files)
        that never go through the normal LLM batch-pagination flow.
        """
        with self._lock:
            if idx := self._indexes.get(scan_id):
                idx.files_served.update(file_paths)
                idx.files_acknowledged.update(file_paths)
                # Intentionally do NOT update last_served_batch — that is reserved
                # for LLM-paged batches so acknowledge_last_batch() stays correct.

    def get_file_coverage(self, scan_id: str, all_files: list[str]) -> dict:
        """Return coverage stats comparing served/acknowledged files vs. the full file list."""
        with self._lock:
            idx = self._indexes.get(scan_id)
            if not idx:
                return {}
            all_set = set(all_files)
            never_served = sorted(all_set - idx.files_served)
            served_not_acked = sorted(idx.files_served - idx.files_acknowledged)
            return {
                "totalFiles": len(all_set),
                "filesServed": len(idx.files_served),
                "filesAcknowledged": len(idx.files_acknowledged),
                "neverSentToLLM": never_served,
                "receivedButNotAnalyzed": served_not_acked,
            }

    def fail_scan(self, scan_id: str) -> None:
        """Mark a scan as failed (e.g. stale / LLM stopped paginating)."""
        with self._lock:
            if scan := self._scans.get(scan_id):
                scan.status = "failed"

    def get_scan_age_seconds(self, scan_id: str) -> float | None:
        """Return how many seconds ago the scan was created, or None if not found."""
        with self._lock:
            scan = self._scans.get(scan_id)
            if not scan:
                return None
            try:
                created = datetime.fromisoformat(scan.stored_at)
                return (datetime.now(timezone.utc) - created).total_seconds()
            except Exception:
                return None

    def finalize_scan(self, scan_id: str) -> StoredSecurityScan | None:
        """Finalize a scan: compute summary and mark completed."""
        with self._lock:
            scan = self._scans.get(scan_id)
            if not scan:
                return None
            files_scanned = len(scan.file_list)
            scan.result.summary = self._compute_summary(scan.result.findings, files_scanned)
            scan.status = "completed"
            return scan

    def delete(self, scan_id: str) -> bool:
        """Delete a scan by ID. Returns True if it existed."""
        with self._lock:
            if scan_id not in self._scans:
                return False
            del self._scans[scan_id]
            self._meta.pop(scan_id, None)
            self._indexes.pop(scan_id, None)
            return True

    def clear(self) -> int:
        """Delete all scans. Returns the number deleted."""
        with self._lock:
            count = len(self._scans)
            self._scans.clear()
            self._meta.clear()
            self._indexes.clear()
            return count

    def get_recent(self, n: int) -> list[StoredSecurityScan]:
        """Return up to N most recent scans, newest first."""
        with self._lock:
            scans = list(self._scans.values())
        scans.sort(key=lambda s: s.stored_at, reverse=True)
        return scans[:n]

    def mark_pushed(self, scan_id: str) -> None:
        """Mark a scan as pushed to Maeris."""
        with self._lock:
            scan = self._scans.get(scan_id)
            if scan:
                scan.pushed_to_maeris = True
                scan.pushed_at = datetime.now(timezone.utc).isoformat()

    def _enforce_max_size(self) -> None:
        if len(self._scans) < self._max_scans:
            return
        oldest_id = min(self._scans.keys(), key=lambda k: self._scans[k].stored_at)
        del self._scans[oldest_id]
        self._meta.pop(oldest_id, None)
        self._indexes.pop(oldest_id, None)

    @staticmethod
    def _fingerprint(finding: SecurityFinding) -> str:
        # Fingerprint is intentionally rule_id-INDEPENDENT.
        # The LLM generates a new rule_id string on every run for the same issue
        # (e.g. "hardcoded-jwt-secret" vs "owasp_a02_hardcoded_jwt_secret").
        # Using rule_id in the fingerprint caused the same vulnerability to be
        # stored multiple times under different names, inflating the count each run.
        # Deduplication is based on location + content only:
        #   file_path + line_start + snippet_hash (if snippet present)
        if finding.code_snippet:
            normalized = finding.code_snippet.strip().lower()
            snippet_hash = hashlib.md5(normalized.encode()).hexdigest()[:8]
            return f"{finding.file_path}|{finding.line_start}|{snippet_hash}"
        return f"{finding.file_path}|{finding.line_start}"

    @staticmethod
    def _index_finding(indexes: ScanIndexes, finding: SecurityFinding) -> None:
        indexes.by_id[finding.id] = finding
        indexes.by_file.setdefault(finding.file_path, []).append(finding.id)
        indexes.by_rule.setdefault(finding.rule_id, []).append(finding.id)
        if finding.owasp_category:
            indexes.by_owasp.setdefault(finding.owasp_category.value, []).append(finding.id)
        indexes.by_severity.setdefault(finding.severity.value, []).append(finding.id)

    @staticmethod
    def _compute_summary(findings: list[SecurityFinding], files_scanned: int) -> ScanSummary:
        severity_counts = {"low": 0, "medium": 0, "high": 0, "critical": 0}
        owasp_counts: dict[str, int] = {}
        for finding in findings:
            sev = finding.severity.value if finding.severity else "medium"
            severity_counts[sev] = severity_counts.get(sev, 0) + 1
            if finding.owasp_category:
                code = finding.owasp_category.value
                owasp_counts[code] = owasp_counts.get(code, 0) + 1

        return ScanSummary(
            total_findings=len(findings),
            by_severity=SeverityCounts(**severity_counts),
            by_owasp_category=owasp_counts,
            files_scanned=files_scanned,
        )


# ---------------------------------------------------------------------------
# Persistent findings cache — module-level helpers (no SecurityScanStore state)
# ---------------------------------------------------------------------------

def _get_findings_cache_path(scan_root: str) -> Path:
    """Return the path to the persistent findings JSON for this scan root."""
    root_hash = hashlib.sha256(scan_root.encode()).hexdigest()[:16]
    base = Path(os.getenv("MAERIS_CONFIG_DIR") or Path.home() / ".maeris")
    cache_dir = base / "security_scans" / root_hash
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir / "findings.json"


def load_persistent_findings(scan_root: str, file_list: list[str], base_path: str) -> list[SecurityFinding]:
    """Load and validate persisted findings for a scan root.

    Drops findings whose source file was deleted or whose code snippet no longer
    exists in the file (meaning the user fixed the vulnerability).

    Each file is read at most once regardless of how many findings reference it,
    so this scales to large codebases with hundreds of cached findings.
    """
    cache_path = _get_findings_cache_path(scan_root)
    if not cache_path.exists():
        return []
    try:
        data = json.loads(cache_path.read_text())
        # Support both old format (bare list) and new format {"findings": [...], "file_hashes": {...}}
        raw_list = data if isinstance(data, list) else data.get("findings", [])
        # model_validate handles camelCase aliases correctly (saved with by_alias=True)
        findings = [SecurityFinding.model_validate(f) for f in raw_list]
    except Exception:
        return []

    file_set = set(file_list)

    # Group findings by file so each file is read from disk exactly once.
    by_file: dict[str, list[SecurityFinding]] = {}
    for finding in findings:
        if finding.file_path in file_set:
            by_file.setdefault(finding.file_path, []).append(finding)

    valid: list[SecurityFinding] = []
    for rel_path, file_findings in by_file.items():
        # Separate findings that need snippet validation from those that don't.
        needs_check = [f for f in file_findings if f.code_snippet]
        no_check = [f for f in file_findings if not f.code_snippet]

        valid.extend(no_check)  # no snippet → keep unconditionally

        if needs_check:
            try:
                content = (Path(base_path) / rel_path).read_text(errors="replace")
            except OSError:
                continue  # can't read file → skip all its findings
            normalized_content = re.sub(r"\s+", " ", content)
            content_lines = [ln.strip() for ln in content.splitlines()]
            for finding in needs_check:
                snippet = finding.code_snippet  # type: ignore[union-attr]
                # 1. Exact match
                exact = snippet.strip() in content
                # 2. Whitespace-normalized match (handles indent/line-ending drift)
                fuzzy = re.sub(r"\s+", " ", snippet.strip()) in normalized_content
                # 3. Anchor match — at least one non-trivial line from the snippet
                #    exists verbatim in the file (handles LLM adding/removing
                #    surrounding context lines between runs).
                anchor = False
                if not exact and not fuzzy:
                    snippet_lines = [
                        ln.strip()
                        for ln in snippet.splitlines()
                        if len(ln.strip()) > 10  # skip blank/trivial lines
                    ]
                    if snippet_lines:
                        anchor = any(ln in content_lines for ln in snippet_lines)
                if exact or fuzzy or anchor:
                    valid.append(finding)
                # else: snippet genuinely gone → user fixed it → drop

    return valid


def load_file_hashes(scan_root: str) -> dict[str, str]:
    """Load saved file content hashes from the previous scan cache.

    Returns an empty dict if no cache exists or the cache is in the old list-only format.
    """
    cache_path = _get_findings_cache_path(scan_root)
    if not cache_path.exists():
        return {}
    try:
        data = json.loads(cache_path.read_text())
        if isinstance(data, dict):
            return data.get("file_hashes", {})
    except Exception:
        pass
    return {}


def save_scan_cache(
    scan_root: str,
    findings: list[SecurityFinding],
    file_hashes: dict[str, str],
) -> None:
    """Persist findings + file content hashes to disk (best-effort).

    The combined cache enables two features on subsequent scans:
    1. Persistent findings — unfixed vulnerabilities survive across runs.
    2. Change-aware scanning — unchanged files (same hash) skip full LLM re-analysis.
    """
    cache_path = _get_findings_cache_path(scan_root)
    try:
        cache_path.write_text(
            json.dumps(
                {
                    "findings": [f.model_dump(by_alias=True) for f in findings],
                    "file_hashes": file_hashes,
                },
                indent=2,
            )
        )
    except Exception:
        pass  # never crash the scan over a cache write failure
