from __future__ import annotations

import json
import sqlite3
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Iterator
from uuid import uuid4

from .models import (
    BusinessResume,
    ComponentBlock,
    ConversationTurn,
    SessionView,
)


def utc_now() -> str:
    return datetime.now(UTC).isoformat()


class Database:
    def __init__(self, path: str | Path) -> None:
        self.path = str(path)
        if self.path != ":memory:":
            Path(self.path).parent.mkdir(parents=True, exist_ok=True)

    def connect(self) -> sqlite3.Connection:
        connection = sqlite3.connect(self.path, timeout=10, isolation_level=None)
        connection.row_factory = sqlite3.Row
        connection.execute("PRAGMA foreign_keys = ON")
        connection.execute("PRAGMA busy_timeout = 10000")
        if self.path != ":memory:":
            connection.execute("PRAGMA journal_mode = WAL")
        return connection

    @contextmanager
    def transaction(self, *, immediate: bool = False) -> Iterator[sqlite3.Connection]:
        connection = self.connect()
        try:
            connection.execute("BEGIN IMMEDIATE" if immediate else "BEGIN")
            yield connection
            connection.commit()
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()

    def initialize(self) -> None:
        with self.transaction(immediate=True) as connection:
            connection.executescript(
                """
                CREATE TABLE IF NOT EXISTS sessions (
                    id TEXT PRIMARY KEY,
                    stage TEXT NOT NULL,
                    revision INTEGER NOT NULL DEFAULT 0,
                    profile_json TEXT NOT NULL,
                    draft_id TEXT,
                    resume_id TEXT,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );

                CREATE TABLE IF NOT EXISTS turns (
                    id TEXT PRIMARY KEY,
                    session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
                    sequence INTEGER NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT,
                    composer_mode TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    UNIQUE(session_id, sequence)
                );

                CREATE TABLE IF NOT EXISTS blocks (
                    id TEXT PRIMARY KEY,
                    session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
                    turn_id TEXT NOT NULL REFERENCES turns(id) ON DELETE CASCADE,
                    block_index INTEGER NOT NULL,
                    type TEXT NOT NULL,
                    lifecycle TEXT NOT NULL,
                    data_json TEXT NOT NULL,
                    version INTEGER NOT NULL DEFAULT 1,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    UNIQUE(turn_id, block_index)
                );

                CREATE TABLE IF NOT EXISTS resumes (
                    id TEXT PRIMARY KEY,
                    session_id TEXT NOT NULL UNIQUE REFERENCES sessions(id) ON DELETE CASCADE,
                    idempotency_key TEXT,
                    revision INTEGER NOT NULL,
                    content_json TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );

                CREATE INDEX IF NOT EXISTS idx_turns_session
                    ON turns(session_id, sequence);
                CREATE INDEX IF NOT EXISTS idx_blocks_session
                    ON blocks(session_id, turn_id, block_index);
                """
            )

    def create_session(
        self,
        session_id: str,
        stage: str,
        profile: dict[str, Any],
        initial_turn: dict[str, Any],
    ) -> None:
        now = utc_now()
        with self.transaction(immediate=True) as connection:
            connection.execute(
                """INSERT INTO sessions
                   (id, stage, revision, profile_json, created_at, updated_at)
                   VALUES (?, ?, 0, ?, ?, ?)""",
                (session_id, stage, json.dumps(profile, ensure_ascii=False), now, now),
            )
            self.insert_turn(connection, session_id=session_id, **initial_turn)

    def fetch_session(
        self, connection: sqlite3.Connection, session_id: str
    ) -> dict[str, Any] | None:
        row = connection.execute(
            "SELECT * FROM sessions WHERE id = ?", (session_id,)
        ).fetchone()
        if row is None:
            return None
        result = dict(row)
        result["profile"] = json.loads(result.pop("profile_json"))
        return result

    def get_session(self, session_id: str) -> dict[str, Any] | None:
        with self.transaction() as connection:
            return self.fetch_session(connection, session_id)

    def update_session(
        self,
        connection: sqlite3.Connection,
        session_id: str,
        *,
        stage: str,
        profile: dict[str, Any],
        draft_id: str | None = None,
        resume_id: str | None = None,
        increment_revision: bool = True,
    ) -> dict[str, Any]:
        current = self.fetch_session(connection, session_id)
        if current is None:
            raise KeyError(session_id)
        revision = current["revision"] + (1 if increment_revision else 0)
        draft_value = draft_id if draft_id is not None else current["draft_id"]
        resume_value = resume_id if resume_id is not None else current["resume_id"]
        connection.execute(
            """UPDATE sessions
               SET stage = ?, revision = ?, profile_json = ?, draft_id = ?,
                   resume_id = ?, updated_at = ?
               WHERE id = ?""",
            (
                stage,
                revision,
                json.dumps(profile, ensure_ascii=False),
                draft_value,
                resume_value,
                utc_now(),
                session_id,
            ),
        )
        updated = self.fetch_session(connection, session_id)
        assert updated is not None
        return updated

    def insert_turn(
        self,
        connection: sqlite3.Connection,
        *,
        session_id: str,
        role: str,
        content: str | None,
        composer_mode: str,
        blocks: list[dict[str, Any]],
    ) -> str:
        turn_id = f"turn_{uuid4().hex}"
        sequence = connection.execute(
            "SELECT COALESCE(MAX(sequence), 0) + 1 FROM turns WHERE session_id = ?",
            (session_id,),
        ).fetchone()[0]
        now = utc_now()
        connection.execute(
            """INSERT INTO turns
               (id, session_id, sequence, role, content, composer_mode, created_at)
               VALUES (?, ?, ?, ?, ?, ?, ?)""",
            (turn_id, session_id, sequence, role, content, composer_mode, now),
        )
        for index, block in enumerate(blocks):
            connection.execute(
                """INSERT INTO blocks
                   (id, session_id, turn_id, block_index, type, lifecycle,
                    data_json, version, created_at, updated_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?, 1, ?, ?)""",
                (
                    block.get("id", f"block_{uuid4().hex}"),
                    session_id,
                    turn_id,
                    index,
                    block["type"],
                    block.get("lifecycle", "active"),
                    json.dumps(block.get("data", {}), ensure_ascii=False),
                    now,
                    now,
                ),
            )
        return turn_id

    def fetch_block(
        self, connection: sqlite3.Connection, session_id: str, block_id: str
    ) -> dict[str, Any] | None:
        row = connection.execute(
            "SELECT * FROM blocks WHERE id = ? AND session_id = ?",
            (block_id, session_id),
        ).fetchone()
        if row is None:
            return None
        result = dict(row)
        result["data"] = json.loads(result.pop("data_json"))
        return result

    def update_block(
        self,
        connection: sqlite3.Connection,
        block_id: str,
        *,
        lifecycle: str,
        data: dict[str, Any] | None = None,
    ) -> None:
        row = connection.execute(
            "SELECT data_json FROM blocks WHERE id = ?", (block_id,)
        ).fetchone()
        if row is None:
            raise KeyError(block_id)
        serialized = row["data_json"] if data is None else json.dumps(data, ensure_ascii=False)
        connection.execute(
            """UPDATE blocks
               SET lifecycle = ?, data_json = ?, version = version + 1, updated_at = ?
               WHERE id = ?""",
            (lifecycle, serialized, utc_now(), block_id),
        )

    def supersede_active_components(
        self,
        connection: sqlite3.Connection,
        session_id: str,
    ) -> None:
        """Make component submissions single-use when chat or create advances the flow."""
        connection.execute(
            """UPDATE blocks
               SET lifecycle = 'superseded', version = version + 1, updated_at = ?
               WHERE session_id = ? AND type = 'component' AND lifecycle = 'active'""",
            (utc_now(), session_id),
        )

    def get_turn(self, turn_id: str) -> ConversationTurn:
        with self.transaction() as connection:
            return self.fetch_turn(connection, turn_id)

    def fetch_turn(
        self,
        connection: sqlite3.Connection,
        turn_id: str,
    ) -> ConversationTurn:
        row = connection.execute("SELECT * FROM turns WHERE id = ?", (turn_id,)).fetchone()
        if row is None:
            raise KeyError(turn_id)
        return self._turn_from_row(connection, row)

    def list_turns(self, session_id: str) -> list[ConversationTurn]:
        with self.transaction() as connection:
            rows = connection.execute(
                "SELECT * FROM turns WHERE session_id = ? ORDER BY sequence",
                (session_id,),
            ).fetchall()
            return [self._turn_from_row(connection, row) for row in rows]

    def _turn_from_row(
        self, connection: sqlite3.Connection, row: sqlite3.Row
    ) -> ConversationTurn:
        block_rows = connection.execute(
            "SELECT * FROM blocks WHERE turn_id = ? ORDER BY block_index", (row["id"],)
        ).fetchall()
        blocks = [
            ComponentBlock(
                id=block["id"],
                type=block["type"],
                lifecycle=block["lifecycle"],
                data=json.loads(block["data_json"]),
                version=block["version"],
                created_at=block["created_at"],
                updated_at=block["updated_at"],
            )
            for block in block_rows
        ]
        return ConversationTurn(
            id=row["id"],
            sequence=row["sequence"],
            role=row["role"],
            content=row["content"],
            composer_mode=row["composer_mode"],
            blocks=blocks,
            created_at=row["created_at"],
        )

    def session_view(self, session: dict[str, Any]) -> SessionView:
        profile = session["profile"]
        phone = profile.get("phone")
        masked_phone = f"{phone[:3]}****{phone[-4:]}" if phone else None
        return SessionView(
            id=session["id"],
            stage=session["stage"],
            revision=session["revision"],
            job_type=profile.get("job_type"),
            anchor_type=profile.get("anchor_type"),
            masked_phone=masked_phone,
            phone_source=profile.get("phone_source"),
            name=profile.get("name"),
            draft_id=session.get("draft_id"),
            resume_id=session.get("resume_id"),
            created_at=session["created_at"],
            updated_at=session["updated_at"],
        )

    def fetch_resume(
        self, connection: sqlite3.Connection, session_id: str
    ) -> dict[str, Any] | None:
        row = connection.execute(
            "SELECT * FROM resumes WHERE session_id = ?", (session_id,)
        ).fetchone()
        if row is None:
            return None
        result = dict(row)
        result["content"] = json.loads(result.pop("content_json"))
        return result

    def insert_resume(
        self,
        connection: sqlite3.Connection,
        *,
        resume_id: str,
        session_id: str,
        idempotency_key: str | None,
        content: dict[str, Any],
    ) -> dict[str, Any]:
        now = utc_now()
        connection.execute(
            """INSERT INTO resumes
               (id, session_id, idempotency_key, revision, content_json, created_at, updated_at)
               VALUES (?, ?, ?, 1, ?, ?, ?)""",
            (
                resume_id,
                session_id,
                idempotency_key,
                json.dumps(content, ensure_ascii=False),
                now,
                now,
            ),
        )
        result = self.fetch_resume(connection, session_id)
        assert result is not None
        return result

    def update_resume(
        self,
        connection: sqlite3.Connection,
        session_id: str,
        content: dict[str, Any],
    ) -> dict[str, Any]:
        connection.execute(
            """UPDATE resumes
               SET revision = revision + 1, content_json = ?, updated_at = ?
               WHERE session_id = ?""",
            (json.dumps(content, ensure_ascii=False), utc_now(), session_id),
        )
        result = self.fetch_resume(connection, session_id)
        if result is None:
            raise KeyError(session_id)
        return result

    @staticmethod
    def resume_view(resume: dict[str, Any]) -> BusinessResume:
        return BusinessResume(
            id=resume["id"],
            session_id=resume["session_id"],
            revision=resume["revision"],
            content=resume["content"],
            created_at=resume["created_at"],
            updated_at=resume["updated_at"],
        )

    def delete_session(self, session_id: str) -> bool:
        with self.transaction(immediate=True) as connection:
            cursor = connection.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
            return cursor.rowcount > 0