"""Audit store: append-only, hash-chained, tamper-evident record storage.

The store guarantees:
  - Every record is written synchronously before the proxy returns a response
  - Each record's chain_hash links to the previous record's chain_hash
  - The chain can be verified at any time to detect tampering
  - Records are never modified or deleted (append-only)
"""

from __future__ import annotations

import sqlite3
import threading
from typing import TYPE_CHECKING, Protocol

from attestmcp.canonical import (
    GENESIS_HASH,
    canonical_json,
    compute_chain_hash,
)
from attestmcp.models import AuditRecord, ChainVerification

if TYPE_CHECKING:
    from datetime import datetime
    from pathlib import Path


class AuditStore(Protocol):
    """Protocol for audit record storage backends.

    Implementations must guarantee:
    1. append() is synchronous and durable before returning
    2. Hash chain integrity is maintained across all records
    3. Records are never modified after append
    """

    def append(self, record: AuditRecord) -> AuditRecord:
        """Append a record to the store, computing and setting its chain_hash.

        The record's chain_hash field will be set by this method.
        All other fields must already be populated.

        Returns the record with chain_hash set.
        Raises RuntimeError if the write fails.
        """
        ...

    def query(
        self,
        since: datetime | None = None,
        until: datetime | None = None,
        limit: int = 1000,
    ) -> list[AuditRecord]:
        """Query records within a time range."""
        ...

    def verify_chain(
        self,
        since: datetime | None = None,
        until: datetime | None = None,
    ) -> ChainVerification:
        """Verify the integrity of the hash chain over a range of records."""
        ...

    def get_chain_head(self) -> str:
        """Return the chain_hash of the most recent record, or GENESIS_HASH if empty."""
        ...

    def count(self) -> int:
        """Return the total number of records in the store."""
        ...


class SQLiteAuditStore:
    """SQLite-backed audit store with WAL mode and synchronous writes.

    Thread-safe: uses a lock around all write operations to ensure
    hash chain integrity even under concurrent access.

    The database schema is intentionally simple:
    - One table, one row per audit record
    - chain_hash links each row to the previous
    - No indexes beyond the primary key (add as needed for query patterns)
    """

    def __init__(self, db_path: str | Path = "attestmcp_audit.db") -> None:
        self._db_path = str(db_path)
        self._write_lock = threading.Lock()
        self._init_db()

    def _get_conn(self) -> sqlite3.Connection:
        """Create a new connection (SQLite connections are not thread-safe)."""
        conn = sqlite3.connect(self._db_path)
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA synchronous=NORMAL")
        conn.execute("PRAGMA busy_timeout=5000")
        conn.row_factory = sqlite3.Row
        return conn

    def _init_db(self) -> None:
        """Create the audit table if it doesn't exist."""
        conn = self._get_conn()
        try:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS audit_records (
                    record_id TEXT PRIMARY KEY,
                    timestamp TEXT NOT NULL,
                    session_id TEXT,
                    agent_identity_json TEXT NOT NULL,
                    method TEXT NOT NULL,
                    tool_name TEXT,
                    params_hash TEXT NOT NULL,
                    policy_decision TEXT NOT NULL,
                    policy_rule TEXT,
                    result_hash TEXT,
                    duration_ms INTEGER NOT NULL DEFAULT 0,
                    error TEXT,
                    chain_hash TEXT NOT NULL,
                    record_json TEXT NOT NULL,
                    sequence_num INTEGER NOT NULL
                )
            """)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_audit_timestamp
                ON audit_records(timestamp)
            """)
            conn.commit()
        finally:
            conn.close()

    def append(self, record: AuditRecord) -> AuditRecord:
        """Append a record with computed chain_hash. Synchronous and durable."""
        with self._write_lock:
            conn = self._get_conn()
            try:
                # Get the current chain head
                row = conn.execute(
                    "SELECT chain_hash, sequence_num FROM audit_records "
                    "ORDER BY sequence_num DESC LIMIT 1"
                ).fetchone()

                if row is None:
                    previous_hash = GENESIS_HASH
                    sequence_num = 0
                else:
                    previous_hash = row["chain_hash"]
                    sequence_num = row["sequence_num"] + 1

                # Compute canonical JSON (excluding chain_hash) for hashing
                record_dict = record.model_dump(mode="python")
                record_canonical = canonical_json(
                    record_dict, exclude_fields={"chain_hash"}
                )

                # Compute chain hash
                chain_hash = compute_chain_hash(previous_hash, record_canonical)
                record.chain_hash = chain_hash

                # Serialize the full record including chain_hash
                record_json = record.model_dump_json()
                identity_json = record.agent_identity.model_dump_json()

                # Insert — synchronous, durable
                conn.execute(
                    """
                    INSERT INTO audit_records (
                        record_id, timestamp, session_id,
                        agent_identity_json, method, tool_name,
                        params_hash, policy_decision, policy_rule,
                        result_hash, duration_ms, error,
                        chain_hash, record_json, sequence_num
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        record.record_id,
                        record.timestamp.isoformat(),
                        record.session_id,
                        identity_json,
                        record.method,
                        record.tool_name,
                        record.params_hash,
                        record.policy_decision.value,
                        record.policy_rule,
                        record.result_hash,
                        record.duration_ms,
                        record.error,
                        chain_hash,
                        record_json,
                        sequence_num,
                    ),
                )
                conn.commit()
                return record

            except Exception:
                conn.rollback()
                raise
            finally:
                conn.close()

    def query(
        self,
        since: datetime | None = None,
        until: datetime | None = None,
        limit: int = 1000,
    ) -> list[AuditRecord]:
        """Query records within a time range, ordered by sequence."""
        conn = self._get_conn()
        try:
            conditions: list[str] = []
            params: list[str | int] = []

            if since is not None:
                conditions.append("timestamp >= ?")
                params.append(since.isoformat())
            if until is not None:
                conditions.append("timestamp <= ?")
                params.append(until.isoformat())

            where = f"WHERE {' AND '.join(conditions)}" if conditions else ""
            sql = f"""
                SELECT record_json FROM audit_records
                {where}
                ORDER BY sequence_num ASC
                LIMIT ?
            """  # noqa: S608 — where clause built from hardcoded columns + parameterized values
            params.append(limit)

            rows = conn.execute(sql, params).fetchall()
            return [AuditRecord.model_validate_json(row["record_json"]) for row in rows]
        finally:
            conn.close()

    def verify_chain(
        self,
        since: datetime | None = None,
        until: datetime | None = None,
    ) -> ChainVerification:
        """Verify hash chain integrity across a range of records.

        Recomputes every chain_hash from scratch and compares
        against the stored value. Any mismatch indicates tampering.
        """
        conn = self._get_conn()
        try:
            conditions: list[str] = []
            params: list[str] = []

            if since is not None:
                conditions.append("timestamp >= ?")
                params.append(since.isoformat())
            if until is not None:
                conditions.append("timestamp <= ?")
                params.append(until.isoformat())

            where = f"WHERE {' AND '.join(conditions)}" if conditions else ""
            sql = f"""
                SELECT record_json, chain_hash, sequence_num
                FROM audit_records
                {where}
                ORDER BY sequence_num ASC
            """  # noqa: S608 — where clause built from hardcoded columns + parameterized values
            rows = conn.execute(sql, params).fetchall()

            if not rows:
                return ChainVerification.empty()

            # For the first record in our range, we need the previous hash
            first_seq = rows[0]["sequence_num"]
            if first_seq == 0:
                previous_hash = GENESIS_HASH
            else:
                prev_row = conn.execute(
                    "SELECT chain_hash FROM audit_records WHERE sequence_num = ?",
                    (first_seq - 1,),
                ).fetchone()
                if prev_row is None:
                    msg = f"Missing record at sequence {first_seq - 1}"
                    raise RuntimeError(msg)
                previous_hash = prev_row["chain_hash"]

            first_record_id: str | None = None
            last_record_id: str | None = None

            for row in rows:
                record = AuditRecord.model_validate_json(row["record_json"])

                if first_record_id is None:
                    first_record_id = record.record_id
                last_record_id = record.record_id

                # Recompute the chain hash
                record_dict = record.model_dump(mode="python")
                record_canonical = canonical_json(
                    record_dict, exclude_fields={"chain_hash"}
                )
                expected_hash = compute_chain_hash(previous_hash, record_canonical)

                stored_hash = row["chain_hash"]
                if expected_hash != stored_hash:
                    return ChainVerification(
                        valid=False,
                        records_checked=rows.index(row) + 1,
                        first_record_id=first_record_id,
                        last_record_id=record.record_id,
                        break_at_record_id=record.record_id,
                        expected_hash=expected_hash,
                        actual_hash=stored_hash,
                    )

                previous_hash = stored_hash

            return ChainVerification(
                valid=True,
                records_checked=len(rows),
                first_record_id=first_record_id,
                last_record_id=last_record_id,
            )
        finally:
            conn.close()

    def get_chain_head(self) -> str:
        """Return the chain_hash of the most recent record."""
        conn = self._get_conn()
        try:
            row = conn.execute(
                "SELECT chain_hash FROM audit_records ORDER BY sequence_num DESC LIMIT 1"
            ).fetchone()
            return row["chain_hash"] if row else GENESIS_HASH
        finally:
            conn.close()

    def count(self) -> int:
        """Return total number of audit records."""
        conn = self._get_conn()
        try:
            row = conn.execute("SELECT COUNT(*) as cnt FROM audit_records").fetchone()
            return row["cnt"] if row else 0
        finally:
            conn.close()
