"""Security Baseline Manager for MCP Audit.

Advanced baseline management with temporal tracking, drift detection,
and automatic evolution. Extends the simple baseline.py with richer
comparison, diffing, and auto-update capabilities.

Zero external dependencies -- pure Python stdlib.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any


# ---------- Data Models ----------

@dataclass
class BaselineFinding:
    """A single finding stored in a baseline."""

    fingerprint: str
    rule_id: str
    severity: str
    title: str
    file: str = ""
    category: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "fingerprint": self.fingerprint,
            "rule_id": self.rule_id,
            "severity": self.severity,
            "title": self.title,
            "file": self.file,
            "category": self.category,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> BaselineFinding:
        return cls(
            fingerprint=d.get("fingerprint", ""),
            rule_id=d.get("rule_id", ""),
            severity=d.get("severity", "MEDIUM"),
            title=d.get("title", ""),
            file=d.get("file", ""),
            category=d.get("category", ""),
        )


@dataclass
class Baseline:
    """A security baseline snapshot."""

    version: int = 2
    created_at: str = ""
    updated_at: str = ""
    label: str = ""
    findings: list[BaselineFinding] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        if not self.created_at:
            self.created_at = datetime.now(timezone.utc).isoformat()
        if not self.updated_at:
            self.updated_at = self.created_at

    @property
    def count(self) -> int:
        return len(self.findings)

    @property
    def fingerprints(self) -> set[str]:
        return {f.fingerprint for f in self.findings}

    @property
    def severity_counts(self) -> dict[str, int]:
        counts: dict[str, int] = {}
        for f in self.findings:
            counts[f.severity] = counts.get(f.severity, 0) + 1
        return counts

    def to_dict(self) -> dict[str, Any]:
        return {
            "version": self.version,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "label": self.label,
            "count": self.count,
            "severity_counts": self.severity_counts,
            "findings": [f.to_dict() for f in self.findings],
            "metadata": self.metadata,
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), indent=2, ensure_ascii=False)

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Baseline:
        findings = [BaselineFinding.from_dict(f) for f in d.get("findings", [])]
        return cls(
            version=d.get("version", 2),
            created_at=d.get("created_at", ""),
            updated_at=d.get("updated_at", ""),
            label=d.get("label", ""),
            findings=findings,
            metadata=d.get("metadata", {}),
        )

    @classmethod
    def from_json(cls, text: str) -> Baseline:
        return cls.from_dict(json.loads(text))


@dataclass
class DiffFinding:
    """A finding in a baseline diff."""

    fingerprint: str
    rule_id: str
    severity: str
    title: str
    file: str = ""
    status: str = ""  # "new", "resolved", "unchanged", "drifted"

    def to_dict(self) -> dict[str, Any]:
        return {
            "fingerprint": self.fingerprint,
            "rule_id": self.rule_id,
            "severity": self.severity,
            "title": self.title,
            "file": self.file,
            "status": self.status,
        }


@dataclass
class BaselineDiff:
    """Difference between a scan result and a baseline."""

    new_findings: list[DiffFinding] = field(default_factory=list)
    resolved_findings: list[DiffFinding] = field(default_factory=list)
    unchanged_findings: list[DiffFinding] = field(default_factory=list)
    drifted_findings: list[DiffFinding] = field(default_factory=list)
    baseline_count: int = 0
    scan_count: int = 0
    drift_score: int = 0  # 0-100

    @property
    def has_changes(self) -> bool:
        return bool(self.new_findings or self.resolved_findings or self.drifted_findings)

    @property
    def regression_count(self) -> int:
        """New findings = regressions."""
        return len(self.new_findings)

    @property
    def improvement_count(self) -> int:
        """Resolved findings = improvements."""
        return len(self.resolved_findings)

    def to_dict(self) -> dict[str, Any]:
        return {
            "new_findings": [f.to_dict() for f in self.new_findings],
            "resolved_findings": [f.to_dict() for f in self.resolved_findings],
            "unchanged_findings": [f.to_dict() for f in self.unchanged_findings],
            "drifted_findings": [f.to_dict() for f in self.drifted_findings],
            "baseline_count": self.baseline_count,
            "scan_count": self.scan_count,
            "drift_score": self.drift_score,
            "has_changes": self.has_changes,
            "regression_count": self.regression_count,
            "improvement_count": self.improvement_count,
        }


@dataclass
class BaselineSnapshot:
    """A point-in-time snapshot for evolution tracking."""

    timestamp: str
    label: str
    finding_count: int
    severity_counts: dict[str, int] = field(default_factory=dict)
    fingerprints: set[str] = field(default_factory=set)

    def to_dict(self) -> dict[str, Any]:
        return {
            "timestamp": self.timestamp,
            "label": self.label,
            "finding_count": self.finding_count,
            "severity_counts": self.severity_counts,
            "fingerprints": sorted(self.fingerprints),
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> BaselineSnapshot:
        return cls(
            timestamp=d.get("timestamp", ""),
            label=d.get("label", ""),
            finding_count=d.get("finding_count", 0),
            severity_counts=d.get("severity_counts", {}),
            fingerprints=set(d.get("fingerprints", [])),
        )


@dataclass
class BaselineEvolution:
    """Time series of baseline changes."""

    snapshots: list[BaselineSnapshot] = field(default_factory=list)

    @property
    def trend(self) -> str:
        """Overall trend: improving, worsening, stable."""
        if len(self.snapshots) < 2:
            return "insufficient_data"
        first = self.snapshots[0].finding_count
        last = self.snapshots[-1].finding_count
        if last < first:
            return "improving"
        if last > first:
            return "worsening"
        return "stable"

    @property
    def total_snapshots(self) -> int:
        return len(self.snapshots)

    def add_snapshot(self, baseline: Baseline, label: str = "") -> BaselineSnapshot:
        """Add a baseline state as a new snapshot."""
        snap = BaselineSnapshot(
            timestamp=datetime.now(timezone.utc).isoformat(),
            label=label or baseline.label,
            finding_count=baseline.count,
            severity_counts=baseline.severity_counts,
            fingerprints=baseline.fingerprints,
        )
        self.snapshots.append(snap)
        return snap

    def to_dict(self) -> dict[str, Any]:
        return {
            "trend": self.trend,
            "total_snapshots": self.total_snapshots,
            "snapshots": [s.to_dict() for s in self.snapshots],
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), indent=2, ensure_ascii=False)

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> BaselineEvolution:
        snapshots = [BaselineSnapshot.from_dict(s) for s in d.get("snapshots", [])]
        return cls(snapshots=snapshots)

    @classmethod
    def from_json(cls, text: str) -> BaselineEvolution:
        return cls.from_dict(json.loads(text))


# ---------- Scan Result Abstraction ----------

@dataclass
class ScanFinding:
    """Minimal finding representation for baseline comparison."""

    fingerprint: str
    rule_id: str
    severity: str
    title: str
    file: str = ""
    category: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "fingerprint": self.fingerprint,
            "rule_id": self.rule_id,
            "severity": self.severity,
            "title": self.title,
            "file": self.file,
            "category": self.category,
        }


@dataclass
class ScanResult:
    """Minimal scan result for baseline operations."""

    findings: list[ScanFinding] = field(default_factory=list)
    timestamp: str = ""
    label: str = ""

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.now(timezone.utc).isoformat()

    @property
    def count(self) -> int:
        return len(self.findings)

    def to_dict(self) -> dict[str, Any]:
        return {
            "findings": [f.to_dict() for f in self.findings],
            "timestamp": self.timestamp,
            "label": self.label,
            "count": self.count,
        }


# ---------- Fingerprint Helpers ----------

def compute_fingerprint(rule_id: str, file: str, code_snippet: str = "") -> str:
    """Compute a stable fingerprint for a finding."""
    import os
    parts = [
        rule_id,
        os.path.basename(file) if file else "",
        code_snippet.strip(),
    ]
    raw = "|".join(parts)
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]


# ---------- Core Functions ----------

def create_baseline(scan_result: ScanResult, label: str = "") -> Baseline:
    """Create a baseline from scan results.

    Args:
        scan_result: Current scan result with findings.
        label: Optional label for the baseline.

    Returns:
        Baseline object with all findings.
    """
    findings = []
    for sf in scan_result.findings:
        findings.append(BaselineFinding(
            fingerprint=sf.fingerprint,
            rule_id=sf.rule_id,
            severity=sf.severity,
            title=sf.title,
            file=sf.file,
            category=sf.category,
        ))

    return Baseline(
        label=label,
        findings=findings,
        metadata={"source_count": scan_result.count},
    )


def compare_to_baseline(
    scan_result: ScanResult,
    baseline: Baseline,
) -> BaselineDiff:
    """Compare scan results against an existing baseline.

    Args:
        scan_result: Current scan result.
        baseline: Existing baseline to compare against.

    Returns:
        BaselineDiff with new, resolved, unchanged, and drifted findings.
    """
    baseline_fps = {f.fingerprint: f for f in baseline.findings}
    scan_fps = {f.fingerprint: f for f in scan_result.findings}

    diff = BaselineDiff(
        baseline_count=baseline.count,
        scan_count=scan_result.count,
    )

    # New findings (in scan but not in baseline)
    for fp, sf in scan_fps.items():
        if fp not in baseline_fps:
            diff.new_findings.append(DiffFinding(
                fingerprint=fp,
                rule_id=sf.rule_id,
                severity=sf.severity,
                title=sf.title,
                file=sf.file,
                status="new",
            ))

    # Resolved findings (in baseline but not in scan)
    for fp, bf in baseline_fps.items():
        if fp not in scan_fps:
            diff.resolved_findings.append(DiffFinding(
                fingerprint=fp,
                rule_id=bf.rule_id,
                severity=bf.severity,
                title=bf.title,
                file=bf.file,
                status="resolved",
            ))

    # Unchanged and drifted (in both)
    for fp in baseline_fps:
        if fp in scan_fps:
            bf = baseline_fps[fp]
            sf = scan_fps[fp]
            if bf.severity != sf.severity:
                diff.drifted_findings.append(DiffFinding(
                    fingerprint=fp,
                    rule_id=sf.rule_id,
                    severity=sf.severity,
                    title=sf.title,
                    file=sf.file,
                    status="drifted",
                ))
            else:
                diff.unchanged_findings.append(DiffFinding(
                    fingerprint=fp,
                    rule_id=sf.rule_id,
                    severity=sf.severity,
                    title=sf.title,
                    file=sf.file,
                    status="unchanged",
                ))

    # Drift score: higher = more changes
    total = max(diff.baseline_count, diff.scan_count, 1)
    changes = len(diff.new_findings) + len(diff.resolved_findings) + len(diff.drifted_findings)
    diff.drift_score = min(int((changes / total) * 100), 100)

    return diff


def update_baseline(
    baseline: Baseline,
    diff: BaselineDiff,
    accept_resolved: bool = True,
    accept_new: bool = True,
) -> Baseline:
    """Update baseline based on a diff (after review).

    Args:
        baseline: Current baseline.
        diff: Diff from compare_to_baseline.
        accept_resolved: Remove resolved findings from baseline.
        accept_new: Add new findings to baseline.

    Returns:
        Updated Baseline.
    """
    # Start from current baseline findings
    fps_to_keep: dict[str, BaselineFinding] = {
        f.fingerprint: f for f in baseline.findings
    }

    # Remove resolved if accepted
    if accept_resolved:
        for rf in diff.resolved_findings:
            fps_to_keep.pop(rf.fingerprint, None)

    # Add new if accepted
    if accept_new:
        for nf in diff.new_findings:
            if nf.fingerprint not in fps_to_keep:
                fps_to_keep[nf.fingerprint] = BaselineFinding(
                    fingerprint=nf.fingerprint,
                    rule_id=nf.rule_id,
                    severity=nf.severity,
                    title=nf.title,
                    file=nf.file,
                )

    # Update drifted severities
    for df in diff.drifted_findings:
        if df.fingerprint in fps_to_keep:
            fps_to_keep[df.fingerprint].severity = df.severity

    updated = Baseline(
        version=baseline.version,
        created_at=baseline.created_at,
        updated_at=datetime.now(timezone.utc).isoformat(),
        label=baseline.label,
        findings=list(fps_to_keep.values()),
        metadata={**baseline.metadata, "last_update_reason": "auto-update after review"},
    )
    return updated


def merge_baselines(base: Baseline, other: Baseline) -> Baseline:
    """Merge two baselines, deduplicating by fingerprint.

    Args:
        base: Primary baseline.
        other: Secondary baseline to merge in.

    Returns:
        Merged Baseline (base takes priority on conflicts).
    """
    fps: dict[str, BaselineFinding] = {}
    for f in other.findings:
        fps[f.fingerprint] = f
    for f in base.findings:
        fps[f.fingerprint] = f  # base wins on conflicts

    return Baseline(
        created_at=base.created_at,
        updated_at=datetime.now(timezone.utc).isoformat(),
        label=f"merged: {base.label} + {other.label}",
        findings=list(fps.values()),
        metadata={"merged_from": [base.label, other.label]},
    )


def filter_baseline(
    baseline: Baseline,
    severity: str | None = None,
    rule_id: str | None = None,
    file_pattern: str | None = None,
) -> Baseline:
    """Filter baseline findings by criteria.

    Args:
        baseline: Baseline to filter.
        severity: Keep only this severity level.
        rule_id: Keep only this rule_id (prefix match).
        file_pattern: Keep only files matching this substring.

    Returns:
        Filtered Baseline.
    """
    filtered = []
    for f in baseline.findings:
        if severity and f.severity != severity:
            continue
        if rule_id and not f.rule_id.startswith(rule_id):
            continue
        if file_pattern and file_pattern not in f.file:
            continue
        filtered.append(f)

    return Baseline(
        version=baseline.version,
        created_at=baseline.created_at,
        updated_at=baseline.updated_at,
        label=f"filtered: {baseline.label}",
        findings=filtered,
        metadata={**baseline.metadata, "filter": {
            "severity": severity, "rule_id": rule_id, "file_pattern": file_pattern,
        }},
    )


# ---------- Formatters ----------

def format_baseline_text(baseline: Baseline) -> str:
    """Format baseline as human-readable text."""
    lines: list[str] = []
    lines.append(f"\n  Security Baseline: {baseline.label or '(unlabeled)'}")
    lines.append(f"  Version: {baseline.version}")
    lines.append(f"  Created: {baseline.created_at}")
    lines.append(f"  Updated: {baseline.updated_at}")
    lines.append(f"  Findings: {baseline.count}")
    if baseline.severity_counts:
        counts_str = ", ".join(f"{k}: {v}" for k, v in sorted(baseline.severity_counts.items()))
        lines.append(f"  Severity: {counts_str}")
    lines.append("")
    for f in baseline.findings[:20]:
        lines.append(f"    [{f.severity}] {f.rule_id}: {f.title} ({f.file})")
    if baseline.count > 20:
        lines.append(f"    ... and {baseline.count - 20} more")
    lines.append("")
    return "\n".join(lines)


def format_baseline_json(baseline: Baseline) -> str:
    """Format baseline as JSON."""
    return baseline.to_json()


def format_diff_text(diff: BaselineDiff) -> str:
    """Format baseline diff as human-readable text."""
    lines: list[str] = []
    lines.append("\n  Baseline Comparison")
    lines.append(f"  Baseline: {diff.baseline_count} findings")
    lines.append(f"  Current:  {diff.scan_count} findings")
    lines.append(f"  Drift Score: {diff.drift_score}/100")
    lines.append("")

    if diff.new_findings:
        lines.append(f"  NEW ({len(diff.new_findings)}):")
        for f in diff.new_findings:
            lines.append(f"    + [{f.severity}] {f.rule_id}: {f.title}")
        lines.append("")

    if diff.resolved_findings:
        lines.append(f"  RESOLVED ({len(diff.resolved_findings)}):")
        for f in diff.resolved_findings:
            lines.append(f"    - [{f.severity}] {f.rule_id}: {f.title}")
        lines.append("")

    if diff.drifted_findings:
        lines.append(f"  DRIFTED ({len(diff.drifted_findings)}):")
        for f in diff.drifted_findings:
            lines.append(f"    ~ [{f.severity}] {f.rule_id}: {f.title}")
        lines.append("")

    if not diff.has_changes:
        lines.append("  No changes detected.")
        lines.append("")

    return "\n".join(lines)


def format_diff_json(diff: BaselineDiff) -> str:
    """Format baseline diff as JSON."""
    return json.dumps(diff.to_dict(), indent=2, ensure_ascii=False)


def format_evolution_text(evolution: BaselineEvolution) -> str:
    """Format baseline evolution as text."""
    lines: list[str] = []
    lines.append(f"\n  Baseline Evolution ({evolution.total_snapshots} snapshots)")
    lines.append(f"  Trend: {evolution.trend}")
    lines.append("")
    for s in evolution.snapshots:
        lines.append(f"    [{s.timestamp}] {s.label}: {s.finding_count} findings")
    lines.append("")
    return "\n".join(lines)


def format_evolution_json(evolution: BaselineEvolution) -> str:
    """Format baseline evolution as JSON."""
    return evolution.to_json()
