from __future__ import annotations

import asyncio
import hashlib
import json
from datetime import datetime
from pathlib import Path
from typing import Any
from uuid import uuid4

from pydantic import BaseModel, Field


class AuditEntry(BaseModel):
    """Hash-chained audit log entry."""

    entry_id: str = Field(default_factory=lambda: str(uuid4()))
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    session_id: str
    event_type: str
    payload: dict[str, Any]
    prev_hash: str
    entry_hash: str


class AuditLog:
    """Append-only audit log with integrity verification."""

    def __init__(self, session_id: str, log_path: str = ".nexus/audit.jsonl") -> None:
        self.session_id = session_id
        self.log_path = Path(log_path)
        self.log_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = asyncio.Lock()
        self._prev_hash = "0" * 64
        if self.log_path.exists() and self.log_path.stat().st_size > 0:
            with self.log_path.open("r", encoding="utf-8") as fh:
                lines = [line.strip() for line in fh if line.strip()]
                if lines:
                    last = json.loads(lines[-1])
                    self._prev_hash = last.get("entry_hash", self._prev_hash)

    async def write(self, event_type: str, payload: dict) -> AuditEntry:
        """Append audit entry and advance hash chain."""
        async with self._lock:
            payload_json = json.dumps(payload, sort_keys=True, separators=(",", ":"))
            entry_hash = hashlib.sha256((payload_json + self._prev_hash).encode("utf-8")).hexdigest()
            entry = AuditEntry(
                session_id=self.session_id,
                event_type=event_type,
                payload=payload,
                prev_hash=self._prev_hash,
                entry_hash=entry_hash,
            )
            with self.log_path.open("a", encoding="utf-8") as fh:
                fh.write(entry.model_dump_json() + "\n")
            self._prev_hash = entry.entry_hash
            return entry

    def verify(self) -> tuple[bool, str]:
        """Verify full hash chain from genesis."""
        prev_hash = "0" * 64
        with self.log_path.open("r", encoding="utf-8") as fh:
            for line in fh:
                if not line.strip():
                    continue
                entry = AuditEntry.model_validate_json(line)
                payload_json = json.dumps(entry.payload, sort_keys=True, separators=(",", ":"))
                expected = hashlib.sha256((payload_json + prev_hash).encode("utf-8")).hexdigest()
                if entry.entry_hash != expected:
                    return False, f"tampered at entry {entry.entry_id}"
                prev_hash = entry.entry_hash
        return True, "ok"

    def query(self, event_type: str | None = None, since: datetime | None = None) -> list[AuditEntry]:
        """Query audit entries by type and timestamp cutoff."""
        out: list[AuditEntry] = []
        with self.log_path.open("r", encoding="utf-8") as fh:
            for line in fh:
                if not line.strip():
                    continue
                entry = AuditEntry.model_validate_json(line)
                if event_type and entry.event_type != event_type:
                    continue
                if since and entry.timestamp < since:
                    continue
                out.append(entry)
        return out
