"""Exemption Manager for MCP Security Audit.

Formal exemption system for security findings. Supports expiration dates,
approval workflows, conditions, and full audit trail. More enterprise-grade
than the annotator module.

Usage:
    from mcpaudit.exemptions import create_exemption, apply_exemptions
    ex = create_exemption("CMD-001", "Mitigated by WAF", "security-team")
    filtered = apply_exemptions(findings, [ex])

CLI: mcpaudit exempt add, exempt list, exempt revoke, exempt audit

Zero external dependencies. Part of mcpaudit v2.3.0.
"""

from __future__ import annotations

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


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

VALID_STATUSES = frozenset({
    "active",
    "expired",
    "revoked",
    "pending",
})

DEFAULT_EXEMPTIONS_FILE = ".mcpaudit-exemptions.json"

# Maximum exemption duration in days (365 = 1 year)
MAX_EXEMPTION_DAYS = 365


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------

@dataclass
class Exemption:
    """A formal security finding exemption."""

    exemption_id: str = ""
    finding_rule: str = ""
    reason: str = ""
    approved_by: str = ""
    status: str = "active"
    created_at: str = ""
    expires_at: str = ""
    revoked_at: str = ""
    revoked_by: str = ""
    conditions: list[str] = field(default_factory=list)
    scope: str = ""  # file pattern, e.g. "server.py" or "*.py"
    notes: str = ""

    def __post_init__(self) -> None:
        now = _now_iso()
        if not self.created_at:
            self.created_at = now
        if not self.exemption_id:
            self.exemption_id = _generate_id(
                self.finding_rule, self.created_at,
            )

    @property
    def is_active(self) -> bool:
        """Check if exemption is currently active (not expired/revoked)."""
        if self.status in ("revoked", "pending"):
            return False
        if self.status == "expired":
            return False
        if self.expires_at:
            return not _is_expired(self.expires_at)
        return self.status == "active"

    def validate(self) -> list[str]:
        """Validate this exemption.

        Returns:
            List of error messages (empty if valid).
        """
        errors: list[str] = []
        if not self.finding_rule:
            errors.append("finding_rule is required")
        if not self.reason:
            errors.append("reason is required")
        if not self.approved_by:
            errors.append("approved_by is required")
        if self.status and self.status not in VALID_STATUSES:
            errors.append(
                f"Invalid status '{self.status}'. "
                f"Valid: {', '.join(sorted(VALID_STATUSES))}"
            )
        if self.expires_at:
            try:
                _parse_iso(self.expires_at)
            except ValueError:
                errors.append(
                    f"Invalid expires_at format: '{self.expires_at}'. "
                    "Use ISO 8601 (YYYY-MM-DDTHH:MM:SSZ)."
                )
        return errors

    def to_dict(self) -> dict[str, Any]:
        return {
            "exemption_id": self.exemption_id,
            "finding_rule": self.finding_rule,
            "reason": self.reason,
            "approved_by": self.approved_by,
            "status": self.status,
            "created_at": self.created_at,
            "expires_at": self.expires_at,
            "revoked_at": self.revoked_at,
            "revoked_by": self.revoked_by,
            "conditions": self.conditions,
            "scope": self.scope,
            "notes": self.notes,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Exemption:
        return cls(
            exemption_id=data.get("exemption_id", ""),
            finding_rule=data.get("finding_rule", ""),
            reason=data.get("reason", ""),
            approved_by=data.get("approved_by", ""),
            status=data.get("status", "active"),
            created_at=data.get("created_at", ""),
            expires_at=data.get("expires_at", ""),
            revoked_at=data.get("revoked_at", ""),
            revoked_by=data.get("revoked_by", ""),
            conditions=data.get("conditions", []),
            scope=data.get("scope", ""),
            notes=data.get("notes", ""),
        )


@dataclass
class AuditEvent:
    """An event in the exemption audit trail."""

    timestamp: str = ""
    action: str = ""  # created, expired, revoked, applied, updated
    exemption_id: str = ""
    finding_rule: str = ""
    actor: str = ""
    details: str = ""

    def __post_init__(self) -> None:
        if not self.timestamp:
            self.timestamp = _now_iso()

    def to_dict(self) -> dict[str, Any]:
        return {
            "timestamp": self.timestamp,
            "action": self.action,
            "exemption_id": self.exemption_id,
            "finding_rule": self.finding_rule,
            "actor": self.actor,
            "details": self.details,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> AuditEvent:
        return cls(
            timestamp=data.get("timestamp", ""),
            action=data.get("action", ""),
            exemption_id=data.get("exemption_id", ""),
            finding_rule=data.get("finding_rule", ""),
            actor=data.get("actor", ""),
            details=data.get("details", ""),
        )


@dataclass
class ExemptionStore:
    """Collection of exemptions with audit trail."""

    exemptions: dict[str, Exemption] = field(default_factory=dict)
    audit_trail: list[AuditEvent] = field(default_factory=list)

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

    @property
    def active_count(self) -> int:
        return sum(1 for e in self.exemptions.values() if e.is_active)

    def get(self, exemption_id: str) -> Exemption | None:
        """Get exemption by ID."""
        return self.exemptions.get(exemption_id)

    def get_by_rule(self, finding_rule: str) -> list[Exemption]:
        """Get all exemptions for a specific rule."""
        return [
            e for e in self.exemptions.values()
            if e.finding_rule == finding_rule
        ]

    def get_active(self) -> list[Exemption]:
        """Get all active exemptions."""
        return [e for e in self.exemptions.values() if e.is_active]

    def get_expired(self) -> list[Exemption]:
        """Get all expired exemptions."""
        result: list[Exemption] = []
        for e in self.exemptions.values():
            if e.expires_at and _is_expired(e.expires_at) and e.status != "revoked":
                result.append(e)
        return result

    def list_all(self) -> list[Exemption]:
        """List all exemptions sorted by creation date."""
        return sorted(
            self.exemptions.values(),
            key=lambda e: e.created_at,
        )

    def stats(self) -> dict[str, int]:
        """Count exemptions by status."""
        counts: dict[str, int] = {}
        for e in self.exemptions.values():
            status = "expired" if (e.expires_at and _is_expired(e.expires_at) and e.status == "active") else e.status
            counts[status] = counts.get(status, 0) + 1
        return counts

    def to_dict(self) -> dict[str, Any]:
        return {
            "exemptions": {
                eid: e.to_dict()
                for eid, e in self.exemptions.items()
            },
            "audit_trail": [a.to_dict() for a in self.audit_trail],
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> ExemptionStore:
        store = cls()
        for eid, edata in data.get("exemptions", {}).items():
            if isinstance(edata, dict):
                ex = Exemption.from_dict(edata)
                if not ex.exemption_id:
                    ex.exemption_id = eid
                store.exemptions[eid] = ex
        for adata in data.get("audit_trail", []):
            if isinstance(adata, dict):
                store.audit_trail.append(AuditEvent.from_dict(adata))
        return store


@dataclass
class FilteredFindings:
    """Result of applying exemptions to findings."""

    included: list[dict[str, Any]] = field(default_factory=list)
    exempted: list[dict[str, Any]] = field(default_factory=list)
    exemption_map: dict[str, str] = field(default_factory=dict)

    @property
    def total_original(self) -> int:
        return len(self.included) + len(self.exempted)

    @property
    def total_exempted(self) -> int:
        return len(self.exempted)

    @property
    def total_remaining(self) -> int:
        return len(self.included)

    def to_dict(self) -> dict[str, Any]:
        return {
            "total_original": self.total_original,
            "total_exempted": self.total_exempted,
            "total_remaining": self.total_remaining,
            "included": self.included,
            "exempted": self.exempted,
            "exemption_map": self.exemption_map,
        }


# ---------------------------------------------------------------------------
# Helper functions
# ---------------------------------------------------------------------------

def _now_iso() -> str:
    """Return current UTC time as ISO string."""
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _parse_iso(iso_str: str) -> datetime:
    """Parse an ISO 8601 datetime string."""
    # Support both Z and +00:00 formats
    s = iso_str.replace("Z", "+00:00")
    if "+" not in s and "-" not in s[10:]:
        s += "+00:00"
    return datetime.fromisoformat(s)


def _is_expired(expires_at: str) -> bool:
    """Check if an expiration date has passed."""
    try:
        expiry = _parse_iso(expires_at)
        now = datetime.now(timezone.utc)
        return now >= expiry
    except (ValueError, TypeError):
        return False


def _generate_id(finding_rule: str, timestamp: str) -> str:
    """Generate a deterministic exemption ID."""
    raw = f"{finding_rule}:{timestamp}"
    return "EX-" + hashlib.sha256(raw.encode()).hexdigest()[:12]


def _matches_scope(scope: str, file_path: str) -> bool:
    """Check if a file path matches an exemption scope pattern."""
    if not scope:
        return True  # No scope means applies globally

    import fnmatch
    # Match against full path and basename
    return (
        fnmatch.fnmatch(file_path, scope)
        or fnmatch.fnmatch(file_path.split("/")[-1], scope)
    )


# ---------------------------------------------------------------------------
# Core operations
# ---------------------------------------------------------------------------

def create_exemption(
    finding_rule: str,
    reason: str,
    approved_by: str,
    expires_at: str = "",
    conditions: list[str] | None = None,
    scope: str = "",
    notes: str = "",
    store: ExemptionStore | None = None,
) -> tuple[Exemption, list[str]]:
    """Create a new exemption.

    Args:
        finding_rule: Rule ID to exempt (e.g., "CMD-001").
        reason: Justification for the exemption.
        approved_by: Who approved this exemption.
        expires_at: Optional expiration date (ISO 8601).
        conditions: Optional conditions for the exemption.
        scope: Optional file pattern scope.
        notes: Optional notes.
        store: Optional ExemptionStore to add to.

    Returns:
        Tuple of (Exemption, errors).
    """
    ex = Exemption(
        finding_rule=finding_rule,
        reason=reason,
        approved_by=approved_by,
        expires_at=expires_at,
        conditions=conditions or [],
        scope=scope,
        notes=notes,
    )

    errors = ex.validate()
    if errors:
        return ex, errors

    if store is not None:
        store.exemptions[ex.exemption_id] = ex
        store.audit_trail.append(AuditEvent(
            action="created",
            exemption_id=ex.exemption_id,
            finding_rule=finding_rule,
            actor=approved_by,
            details=f"Exemption created: {reason}",
        ))

    return ex, []


def revoke_exemption(
    exemption_id: str,
    revoked_by: str,
    store: ExemptionStore,
    reason: str = "",
) -> tuple[bool, str]:
    """Revoke an active exemption.

    Args:
        exemption_id: ID of the exemption to revoke.
        revoked_by: Who is revoking.
        store: ExemptionStore containing the exemption.
        reason: Optional reason for revocation.

    Returns:
        Tuple of (success, message).
    """
    ex = store.get(exemption_id)
    if ex is None:
        return False, f"Exemption not found: {exemption_id}"

    if ex.status == "revoked":
        return False, f"Exemption already revoked: {exemption_id}"

    ex.status = "revoked"
    ex.revoked_at = _now_iso()
    ex.revoked_by = revoked_by

    store.audit_trail.append(AuditEvent(
        action="revoked",
        exemption_id=exemption_id,
        finding_rule=ex.finding_rule,
        actor=revoked_by,
        details=reason or "Exemption revoked",
    ))

    return True, f"Exemption {exemption_id} revoked"


def expire_exemptions(store: ExemptionStore) -> list[str]:
    """Auto-expire all exemptions past their expiration date.

    Args:
        store: ExemptionStore to check.

    Returns:
        List of exemption IDs that were expired.
    """
    expired_ids: list[str] = []
    for eid, ex in store.exemptions.items():
        if (
            ex.status == "active"
            and ex.expires_at
            and _is_expired(ex.expires_at)
        ):
            ex.status = "expired"
            expired_ids.append(eid)
            store.audit_trail.append(AuditEvent(
                action="expired",
                exemption_id=eid,
                finding_rule=ex.finding_rule,
                actor="system",
                details=f"Auto-expired at {ex.expires_at}",
            ))

    return expired_ids


def apply_exemptions(
    findings: list[dict[str, Any]],
    exemptions: list[Exemption] | ExemptionStore,
    auto_expire: bool = True,
) -> FilteredFindings:
    """Apply exemptions to a list of findings.

    Args:
        findings: List of finding dicts.
        exemptions: List of Exemption objects or ExemptionStore.
        auto_expire: Whether to auto-expire before applying.

    Returns:
        FilteredFindings with included and exempted findings.
    """
    if isinstance(exemptions, ExemptionStore):
        if auto_expire:
            expire_exemptions(exemptions)
        active = exemptions.get_active()
    else:
        active = [e for e in exemptions if e.is_active]

    # Build rule -> exemptions lookup
    rule_map: dict[str, list[Exemption]] = {}
    for ex in active:
        rule_map.setdefault(ex.finding_rule, []).append(ex)

    result = FilteredFindings()

    for f in findings:
        rule_id = f.get("rule_id", "")
        file_path = f.get("file", "")

        matched_ex = None
        for ex in rule_map.get(rule_id, []):
            if _matches_scope(ex.scope, file_path):
                matched_ex = ex
                break

        if matched_ex is not None:
            result.exempted.append(f)
            result.exemption_map[f"{rule_id}:{file_path}"] = matched_ex.exemption_id
        else:
            result.included.append(f)

    return result


def update_exemption(
    exemption_id: str,
    store: ExemptionStore,
    reason: str | None = None,
    expires_at: str | None = None,
    conditions: list[str] | None = None,
    notes: str | None = None,
    updated_by: str = "",
) -> tuple[bool, str]:
    """Update an existing exemption.

    Args:
        exemption_id: ID of the exemption to update.
        store: ExemptionStore containing the exemption.
        reason: New reason (or None to keep).
        expires_at: New expiration (or None to keep).
        conditions: New conditions (or None to keep).
        notes: New notes (or None to keep).
        updated_by: Who is updating.

    Returns:
        Tuple of (success, message).
    """
    ex = store.get(exemption_id)
    if ex is None:
        return False, f"Exemption not found: {exemption_id}"

    changes: list[str] = []
    if reason is not None:
        ex.reason = reason
        changes.append("reason")
    if expires_at is not None:
        ex.expires_at = expires_at
        changes.append("expires_at")
    if conditions is not None:
        ex.conditions = conditions
        changes.append("conditions")
    if notes is not None:
        ex.notes = notes
        changes.append("notes")

    if changes:
        store.audit_trail.append(AuditEvent(
            action="updated",
            exemption_id=exemption_id,
            finding_rule=ex.finding_rule,
            actor=updated_by or "unknown",
            details=f"Updated fields: {', '.join(changes)}",
        ))

    return True, f"Exemption {exemption_id} updated: {', '.join(changes)}"


# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------

def load_exemptions(
    path: str | Path = DEFAULT_EXEMPTIONS_FILE,
) -> ExemptionStore:
    """Load exemptions from a JSON file.

    Args:
        path: Path to the exemptions file.

    Returns:
        ExemptionStore (empty if file doesn't exist).
    """
    filepath = Path(path)
    if not filepath.exists():
        return ExemptionStore()

    try:
        text = filepath.read_text(encoding="utf-8")
    except (OSError, UnicodeDecodeError):
        return ExemptionStore()

    try:
        data = json.loads(text)
    except json.JSONDecodeError:
        return ExemptionStore()

    if not isinstance(data, dict):
        return ExemptionStore()

    return ExemptionStore.from_dict(data)


def save_exemptions(
    store: ExemptionStore,
    path: str | Path = DEFAULT_EXEMPTIONS_FILE,
) -> None:
    """Save exemptions to a JSON file.

    Args:
        store: ExemptionStore to persist.
        path: Output file path.
    """
    filepath = Path(path)
    data = store.to_dict()
    filepath.write_text(
        json.dumps(data, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )


def export_exemptions(
    store: ExemptionStore,
) -> str:
    """Export exemptions as JSON string.

    Args:
        store: ExemptionStore to export.

    Returns:
        JSON string.
    """
    return json.dumps(store.to_dict(), indent=2, ensure_ascii=False)


def import_exemptions(
    json_str: str,
    target_store: ExemptionStore | None = None,
) -> tuple[ExemptionStore, int]:
    """Import exemptions from a JSON string.

    Args:
        json_str: JSON string with exemptions data.
        target_store: Optional existing store to merge into.

    Returns:
        Tuple of (ExemptionStore, count_imported).
    """
    try:
        data = json.loads(json_str)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON: {e}")

    if not isinstance(data, dict):
        raise ValueError("Expected JSON object with exemptions key")

    imported = ExemptionStore.from_dict(data)

    if target_store is None:
        return imported, imported.count

    count = 0
    for eid, ex in imported.exemptions.items():
        if eid not in target_store.exemptions:
            target_store.exemptions[eid] = ex
            count += 1
            target_store.audit_trail.append(AuditEvent(
                action="imported",
                exemption_id=eid,
                finding_rule=ex.finding_rule,
                actor="system",
                details="Imported from external source",
            ))

    return target_store, count


# ---------------------------------------------------------------------------
# Audit trail
# ---------------------------------------------------------------------------

def get_audit_trail(
    store: ExemptionStore,
    exemption_id: str | None = None,
    action: str | None = None,
    limit: int = 0,
) -> list[AuditEvent]:
    """Get audit trail events, optionally filtered.

    Args:
        store: ExemptionStore with audit trail.
        exemption_id: Filter by exemption ID.
        action: Filter by action type.
        limit: Max events to return (0 = all).

    Returns:
        List of AuditEvent objects.
    """
    events = store.audit_trail

    if exemption_id:
        events = [e for e in events if e.exemption_id == exemption_id]

    if action:
        events = [e for e in events if e.action == action]

    # Sort by timestamp descending (most recent first)
    events = sorted(events, key=lambda e: e.timestamp, reverse=True)

    if limit > 0:
        events = events[:limit]

    return events


# ---------------------------------------------------------------------------
# Formatting
# ---------------------------------------------------------------------------

def format_exemptions_text(store: ExemptionStore) -> str:
    """Format exemptions as human-readable text.

    Args:
        store: ExemptionStore to format.

    Returns:
        Formatted text.
    """
    if store.count == 0:
        return "No exemptions.\n"

    lines: list[str] = [
        f"Exemptions ({store.count} total, {store.active_count} active)",
        "=" * 60,
    ]

    stats = store.stats()
    for status, count in sorted(stats.items()):
        lines.append(f"  {status}: {count}")
    lines.append("")

    for ex in store.list_all():
        active_str = "ACTIVE" if ex.is_active else ex.status.upper()
        lines.append(f"  [{active_str}] {ex.exemption_id[:16]}...")
        lines.append(f"    Rule: {ex.finding_rule}")
        lines.append(f"    Reason: {ex.reason}")
        lines.append(f"    Approved by: {ex.approved_by}")
        if ex.expires_at:
            lines.append(f"    Expires: {ex.expires_at}")
        if ex.scope:
            lines.append(f"    Scope: {ex.scope}")
        if ex.conditions:
            lines.append(f"    Conditions: {', '.join(ex.conditions)}")
        lines.append(f"    Created: {ex.created_at}")

    return "\n".join(lines) + "\n"


def format_exemptions_json(store: ExemptionStore) -> str:
    """Format exemptions as JSON.

    Args:
        store: ExemptionStore to format.

    Returns:
        JSON string.
    """
    data = {
        "exemptions": {eid: e.to_dict() for eid, e in store.exemptions.items()},
        "count": store.count,
        "active_count": store.active_count,
        "stats": store.stats(),
        "audit_trail": [a.to_dict() for a in store.audit_trail],
    }
    return json.dumps(data, indent=2, ensure_ascii=False)


def format_audit_trail_text(events: list[AuditEvent]) -> str:
    """Format audit trail events as text.

    Args:
        events: List of AuditEvent objects.

    Returns:
        Formatted text.
    """
    if not events:
        return "No audit events.\n"

    lines: list[str] = [
        f"Audit Trail ({len(events)} events)",
        "=" * 60,
    ]

    for e in events:
        lines.append(
            f"  [{e.timestamp}] {e.action.upper()} "
            f"{e.exemption_id[:16]} ({e.finding_rule})"
        )
        if e.actor:
            lines.append(f"    Actor: {e.actor}")
        if e.details:
            lines.append(f"    Details: {e.details}")

    return "\n".join(lines) + "\n"


def format_filtered_text(filtered: FilteredFindings) -> str:
    """Format filtered findings summary as text.

    Args:
        filtered: FilteredFindings result.

    Returns:
        Formatted text.
    """
    lines: list[str] = [
        "Exemption Filter Results",
        "=" * 40,
        f"  Original findings: {filtered.total_original}",
        f"  Exempted: {filtered.total_exempted}",
        f"  Remaining: {filtered.total_remaining}",
    ]

    if filtered.exempted:
        lines.append("")
        lines.append("  Exempted findings:")
        for f in filtered.exempted:
            rule = f.get("rule_id", "?")
            title = f.get("title", "?")
            lines.append(f"    - {rule}: {title}")

    return "\n".join(lines) + "\n"
