from __future__ import annotations

import json
import sqlite3
from pathlib import Path
from typing import Any

from unmake.models import OperationRecord, SessionRecord, SessionStatus, new_id, utcnow
from unmake.utils.helpers import parse_iso_datetime
from unmake.utils.validation import sanitize_op_type, validate_payload


class SQLiteStorage:
    """SQLite-based persistent storage backend.

    Thread Safety:
        This backend is thread-safe in typical usage because each operation
        opens its own SQLite connection and relies on SQLite internal locking.
        Create one HistoryManager per thread for best performance.
    """

    def __init__(self, path: str | Path) -> None:
        self.path = str(path)
        self._init_schema()

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.path)
        conn.row_factory = sqlite3.Row
        return conn

    def _init_schema(self) -> None:
        db_path = Path(self.path)
        if not db_path.exists():
            db_path.touch(mode=0o600)

        with self._connect() as conn:
            conn.executescript(
                """
                PRAGMA foreign_keys = ON;
                PRAGMA journal_mode = WAL;
                PRAGMA synchronous = NORMAL;
                PRAGMA temp_store = MEMORY;
                PRAGMA mmap_size = 30000000000;
                PRAGMA page_size = 4096;

                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    created_at TEXT NOT NULL,
                    completed_at TEXT,
                    status TEXT NOT NULL,
                    context_json TEXT NOT NULL,
                    error TEXT,
                    undo_count INTEGER NOT NULL DEFAULT 0
                );

                CREATE INDEX IF NOT EXISTS idx_sessions_created_at
                    ON sessions(created_at DESC);
                CREATE INDEX IF NOT EXISTS idx_sessions_status
                    ON sessions(status);

                CREATE TABLE IF NOT EXISTS operations (
                    op_id TEXT PRIMARY KEY,
                    session_id TEXT NOT NULL,
                    sequence INTEGER NOT NULL,
                    op_type TEXT NOT NULL,
                    payload_json TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE,
                    UNIQUE (session_id, sequence)
                );

                CREATE INDEX IF NOT EXISTS idx_operations_session_id
                    ON operations(session_id, sequence);
                """
            )

    def start_session(self, *, context: dict[str, Any]) -> SessionRecord:
        record = SessionRecord(
            session_id=new_id(), created_at=utcnow(), status=SessionStatus.RUNNING, context=context
        )
        with self._connect() as conn:
            conn.execute(
                """
                INSERT INTO sessions(session_id, created_at, status, context_json, undo_count)
                VALUES (?, ?, ?, ?, ?)
                """,
                (
                    record.session_id,
                    record.created_at.isoformat(),
                    record.status.value,
                    json.dumps(record.context),
                    record.undo_count,
                ),
            )
        return record

    def append_operation(
        self,
        *,
        session_id: str,
        op_type: str,
        payload: dict[str, Any],
    ) -> OperationRecord:
        op_type = sanitize_op_type(op_type)
        validate_payload(payload)

        with self._connect() as conn:
            exists = conn.execute(
                "SELECT 1 FROM sessions WHERE session_id = ?",
                (session_id,),
            ).fetchone()
            if not exists:
                raise KeyError(
                    f"Session not found: {session_id}. Use list_sessions() to see available sessions."
                )

            row = conn.execute(
                "SELECT COALESCE(MAX(sequence), 0) + 1 AS next_sequence FROM operations WHERE session_id = ?",
                (session_id,),
            ).fetchone()
            seq = int(row["next_sequence"])
            op = OperationRecord(
                op_id=new_id(),
                session_id=session_id,
                sequence=seq,
                op_type=op_type,
                payload=payload,
                created_at=utcnow(),
            )
            conn.execute(
                """
                INSERT INTO operations(op_id, session_id, sequence, op_type, payload_json, created_at)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                (
                    op.op_id,
                    op.session_id,
                    op.sequence,
                    op.op_type,
                    json.dumps(op.payload),
                    op.created_at.isoformat(),
                ),
            )
            return op

    def close_session(
        self,
        *,
        session_id: str,
        status: SessionStatus,
        error: str | None = None,
    ) -> SessionRecord:
        completed_at = utcnow()
        with self._connect() as conn:
            result = conn.execute(
                """
                UPDATE sessions
                SET completed_at = ?, status = ?, error = ?
                WHERE session_id = ?
                """,
                (completed_at.isoformat(), status.value, error, session_id),
            )
            if result.rowcount == 0:
                raise KeyError(
                    f"Session not found: {session_id}. Use list_sessions() to see available sessions."
                )
        session = self.get_session(session_id)
        if session is None:
            raise KeyError(
                f"Session not found: {session_id}. Use list_sessions() to see available sessions."
            )
        return session

    def increment_undo_count(self, session_id: str) -> None:
        with self._connect() as conn:
            conn.execute(
                "UPDATE sessions SET undo_count = undo_count + 1 WHERE session_id = ?",
                (session_id,),
            )

    def list_sessions(self, *, limit: int | None = None) -> list[SessionRecord]:
        query = "SELECT * FROM sessions ORDER BY created_at DESC"
        params: tuple[object, ...] = ()
        if limit:
            query += " LIMIT ?"
            params = (limit,)
        with self._connect() as conn:
            rows = conn.execute(query, params).fetchall()
        return [self._map_session(row) for row in rows]

    def get_session(self, session_id: str) -> SessionRecord | None:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT * FROM sessions WHERE session_id = ?", (session_id,)
            ).fetchone()
        return self._map_session(row) if row else None

    def get_operations(self, session_id: str) -> list[OperationRecord]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT * FROM operations WHERE session_id = ? ORDER BY sequence ASC",
                (session_id,),
            ).fetchall()
        return [self._map_operation(row) for row in rows]

    @staticmethod
    def _map_session(row: sqlite3.Row) -> SessionRecord:
        return SessionRecord(
            session_id=row["session_id"],
            created_at=parse_iso_datetime(row["created_at"]),
            completed_at=parse_iso_datetime(row["completed_at"]) if row["completed_at"] else None,
            status=SessionStatus(row["status"]),
            context=json.loads(row["context_json"]),
            error=row["error"],
            undo_count=int(row["undo_count"]) if "undo_count" in row.keys() else 0,
        )

    @staticmethod
    def _map_operation(row: sqlite3.Row) -> OperationRecord:
        return OperationRecord(
            op_id=row["op_id"],
            session_id=row["session_id"],
            sequence=row["sequence"],
            op_type=row["op_type"],
            payload=json.loads(row["payload_json"]),
            created_at=parse_iso_datetime(row["created_at"]),
        )
