"""Tamper-evident audit logger for AI agent actions.

Logs every LLM call with: timestamp, model, provider, input/output,
token counts, user session, purpose, and PII detection results.

Each entry includes a hash of the previous entry, creating an
audit chain that's tamper-evident — if any entry is modified,
the chain breaks.
"""

import hashlib
import json
import os
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Optional

from . import pii


@dataclass
class AuditEntry:
    """A single audit log entry."""
    timestamp: float
    event_type: str  # "llm_call", "consent", "data_access", "redaction"
    provider: str
    model: str
    purpose: str
    user_id: Optional[str]
    input_text: str
    input_redacted: str
    output_text: str
    pii_detected: dict[str, int]  # PII type → count
    tokens_in: int
    tokens_out: int
    metadata: dict[str, Any] = field(default_factory=dict)
    prev_hash: str = ""
    entry_hash: str = ""


def _hash_entry(entry: dict) -> str:
    """Generate SHA-256 hash of an entry (excluding entry_hash field)."""
    data = {k: v for k, v in entry.items() if k != "entry_hash"}
    raw = json.dumps(data, sort_keys=True, default=str)
    return hashlib.sha256(raw.encode()).hexdigest()[:16]


class AuditLog:
    """Append-only, tamper-evident audit log.

    Stores entries as JSON Lines (.jsonl). Each entry includes a hash
    of the previous entry, forming a verifiable chain.

    Usage:
        log = AuditLog("./audit.jsonl")
        log.record(
            event_type="llm_call",
            provider="openai",
            model="gpt-4",
            purpose="customer_support",
            user_id="user_123",
            input_text="What is my order status?",
            output_text="Your order #456 is being delivered.",
            tokens_in=12,
            tokens_out=15,
        )
    """

    def __init__(self, path: str = "./agent_shield_audit.jsonl"):
        self.path = Path(path)
        self._last_hash = self._get_last_hash()

    def _get_last_hash(self) -> str:
        """Read the hash of the last entry in the log."""
        if not self.path.exists():
            return "genesis"
        try:
            last_line = ""
            with open(self.path, "r") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        last_line = line
            if last_line:
                entry = json.loads(last_line)
                return entry.get("entry_hash", "genesis")
        except (json.JSONDecodeError, OSError):
            pass
        return "genesis"

    def record(
        self,
        event_type: str = "llm_call",
        provider: str = "",
        model: str = "",
        purpose: str = "",
        user_id: Optional[str] = None,
        input_text: str = "",
        output_text: str = "",
        tokens_in: int = 0,
        tokens_out: int = 0,
        metadata: Optional[dict] = None,
    ) -> AuditEntry:
        """Record an event to the audit log.

        PII detection runs automatically on input_text. The redacted
        version is stored alongside the original for compliance review.
        """
        pii_found = pii.summary(input_text)
        input_redacted = pii.redact(input_text) if pii_found else input_text

        entry = AuditEntry(
            timestamp=time.time(),
            event_type=event_type,
            provider=provider,
            model=model,
            purpose=purpose,
            user_id=user_id,
            input_text=input_text,
            input_redacted=input_redacted,
            output_text=output_text,
            pii_detected=pii_found,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            metadata=metadata or {},
            prev_hash=self._last_hash,
        )

        entry_dict = asdict(entry)
        entry.entry_hash = _hash_entry(entry_dict)
        entry_dict["entry_hash"] = entry.entry_hash

        # Append to log file
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.path, "a") as f:
            f.write(json.dumps(entry_dict, default=str) + "\n")

        self._last_hash = entry.entry_hash
        return entry

    def query(
        self,
        user_id: Optional[str] = None,
        event_type: Optional[str] = None,
        since: Optional[float] = None,
    ) -> list[dict]:
        """Query the audit log.

        Supports filtering by user_id (for data subject access requests),
        event_type, and time range.
        """
        if not self.path.exists():
            return []

        results = []
        with open(self.path, "r") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                except json.JSONDecodeError:
                    continue

                if user_id and entry.get("user_id") != user_id:
                    continue
                if event_type and entry.get("event_type") != event_type:
                    continue
                if since and entry.get("timestamp", 0) < since:
                    continue

                results.append(entry)

        return results

    def verify_chain(self) -> tuple[bool, int]:
        """Verify the audit chain integrity.

        Returns:
            (is_valid, entry_count) — True if chain is intact.
        """
        if not self.path.exists():
            return True, 0

        prev_hash = "genesis"
        count = 0

        with open(self.path, "r") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                except json.JSONDecodeError:
                    return False, count

                if entry.get("prev_hash") != prev_hash:
                    return False, count

                expected_hash = _hash_entry(entry)
                if entry.get("entry_hash") != expected_hash:
                    return False, count

                prev_hash = entry["entry_hash"]
                count += 1

        return True, count

    def export_for_user(self, user_id: str) -> list[dict]:
        """Export all data for a specific user (data subject access request).

        Returns entries with sensitive fields included — intended for
        the data subject themselves or the DPO handling the request.
        """
        return self.query(user_id=user_id)

    def stats(self) -> dict:
        """Get summary statistics from the audit log."""
        entries = self.query()
        if not entries:
            return {"total_entries": 0}

        providers = {}
        purposes = {}
        total_pii = {}
        total_tokens_in = 0
        total_tokens_out = 0

        for e in entries:
            p = e.get("provider", "unknown")
            providers[p] = providers.get(p, 0) + 1

            purpose = e.get("purpose", "unknown")
            purposes[purpose] = purposes.get(purpose, 0) + 1

            for pii_type, count in e.get("pii_detected", {}).items():
                total_pii[pii_type] = total_pii.get(pii_type, 0) + count

            total_tokens_in += e.get("tokens_in", 0)
            total_tokens_out += e.get("tokens_out", 0)

        return {
            "total_entries": len(entries),
            "providers": providers,
            "purposes": purposes,
            "pii_detected": total_pii,
            "total_tokens_in": total_tokens_in,
            "total_tokens_out": total_tokens_out,
        }
