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 .resume_document_core import attach_gap_report_staleness
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 TABLE IF NOT EXISTS resume_imports (
                    id TEXT PRIMARY KEY,
                    session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
                    file_name TEXT NOT NULL,
                    mime_type TEXT NOT NULL,
                    size_bytes INTEGER NOT NULL,
                    sha256 TEXT NOT NULL,
                    object_key TEXT NOT NULL,
                    status TEXT NOT NULL,
                    document_json TEXT,
                    field_reviews_json TEXT NOT NULL DEFAULT '[]',
                    error_code TEXT,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    UNIQUE(session_id, sha256)
                );

                CREATE INDEX IF NOT EXISTS idx_resume_imports_session
                    ON resume_imports(session_id, created_at DESC);

                CREATE TABLE IF NOT EXISTS optimization_runs (
                    id TEXT PRIMARY KEY,
                    session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
                    entry_id TEXT NOT NULL,
                    mode TEXT NOT NULL,
                    status TEXT NOT NULL,
                    source_revision INTEGER NOT NULL,
                    state_json TEXT NOT NULL,
                    proposal_json TEXT,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );

                CREATE INDEX IF NOT EXISTS idx_optimization_runs_active
                    ON optimization_runs(session_id, entry_id, status);

                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,
        *,
        for_update: bool = False,
    ) -> dict[str, Any] | None:
        del for_update
        row = connection.execute(
            "SELECT * FROM sessions WHERE id = ?", (session_id,)
        ).fetchone()
        if row is None:
            return None
        result = dict(row)
        profile = json.loads(result.pop("profile_json"))
        if profile.get("job_type") == "other":
            # Sessions created before the workflow upgrade used "other". Treat
            # them as internship sessions so existing users can still resume.
            profile["job_type"] = "internship"
            updated_at = utc_now()
            connection.execute(
                "UPDATE sessions SET profile_json = ?, updated_at = ? WHERE id = ?",
                (json.dumps(profile, ensure_ascii=False), updated_at, session_id),
            )
            result["updated_at"] = updated_at
        result["profile"] = profile
        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 find_latest_session_by_external_user_id(
        self, external_user_id: str
    ) -> dict[str, Any] | None:
        normalized_user_id = str(external_user_id or "").strip()
        if not normalized_user_id:
            return None

        with self.transaction() as connection:
            row = connection.execute(
                """SELECT id
                   FROM sessions
                   WHERE json_extract(
                             profile_json, '$.external_account.provider'
                         ) = 'offerpai'
                     AND CAST(
                             json_extract(
                                 profile_json, '$.external_account.user_id'
                             ) AS TEXT
                         ) = ?
                   ORDER BY updated_at DESC, created_at DESC, id DESC
                   LIMIT 1""",
                (normalized_user_id,),
            ).fetchone()
            return (
                self.fetch_session(connection, row["id"])
                if row is not None
                else None
            )

    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

    def create_optimization_run(
        self,
        connection: sqlite3.Connection,
        *,
        run_id: str,
        session_id: str,
        entry_id: str,
        mode: str,
        status: str,
        source_revision: int,
        state: dict[str, Any],
        proposal: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        now = utc_now()
        connection.execute(
            """INSERT INTO optimization_runs
               (id, session_id, entry_id, mode, status, source_revision, state_json,
                proposal_json, created_at, updated_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                run_id, session_id, entry_id, mode, status, source_revision,
                json.dumps(state, ensure_ascii=False),
                json.dumps(proposal, ensure_ascii=False) if proposal else None, now, now,
            ),
        )
        result = self.fetch_optimization_run(connection, session_id, run_id)
        assert result is not None
        return result

    def fetch_optimization_run(
        self, connection: sqlite3.Connection, session_id: str, run_id: str
    ) -> dict[str, Any] | None:
        row = connection.execute(
            "SELECT * FROM optimization_runs WHERE id = ? AND session_id = ?", (run_id, session_id)
        ).fetchone()
        if row is None:
            return None
        result = dict(row)
        result["state"] = json.loads(result.pop("state_json"))
        raw_proposal = result.pop("proposal_json")
        result["proposal"] = json.loads(raw_proposal) if raw_proposal else None
        return result

    def find_active_optimization_run(
        self, connection: sqlite3.Connection, session_id: str, entry_id: str
    ) -> dict[str, Any] | None:
        row = connection.execute(
            """SELECT id FROM optimization_runs
               WHERE session_id = ? AND entry_id = ?
                 AND status IN ('question_pending', 'proposal_pending')
               ORDER BY created_at DESC LIMIT 1""",
            (session_id, entry_id),
        ).fetchone()
        return self.fetch_optimization_run(connection, session_id, row["id"]) if row else None

    def list_active_optimization_runs(
        self, connection: sqlite3.Connection, session_id: str
    ) -> list[dict[str, Any]]:
        rows = connection.execute(
            """SELECT id FROM optimization_runs
               WHERE session_id = ?
                 AND status IN ('question_pending', 'proposal_pending')
               ORDER BY updated_at ASC, created_at ASC""",
            (session_id,),
        ).fetchall()
        return [
            run for row in rows
            if (run := self.fetch_optimization_run(connection, session_id, row["id"])) is not None
        ]

    def update_optimization_run(
        self,
        connection: sqlite3.Connection,
        *,
        session_id: str,
        run_id: str,
        status: str,
        state: dict[str, Any],
        proposal: dict[str, Any] | None,
    ) -> dict[str, Any]:
        connection.execute(
            """UPDATE optimization_runs
               SET status = ?, state_json = ?, proposal_json = ?, updated_at = ?
               WHERE id = ? AND session_id = ?""",
            (
                status, json.dumps(state, ensure_ascii=False),
                json.dumps(proposal, ensure_ascii=False) if proposal else None,
                utc_now(), run_id, session_id,
            ),
        )
        result = self.fetch_optimization_run(connection, session_id, run_id)
        if result is None:
            raise KeyError(run_id)
        return result
    def create_resume_import(
        self,
        connection: sqlite3.Connection,
        *,
        import_id: str,
        session_id: str,
        file_name: str,
        mime_type: str,
        size_bytes: int,
        sha256: str,
        object_key: str,
        document: dict[str, Any],
        field_reviews: list[dict[str, Any]],
    ) -> dict[str, Any]:
        now = utc_now()
        connection.execute(
            """INSERT INTO resume_imports
               (id, session_id, file_name, mime_type, size_bytes, sha256, object_key,
                status, document_json, field_reviews_json, created_at, updated_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, 'awaiting_review', ?, ?, ?, ?)""",
            (
                import_id, session_id, file_name, mime_type, size_bytes, sha256, object_key,
                json.dumps(document, ensure_ascii=False),
                json.dumps(field_reviews, ensure_ascii=False), now, now,
            ),
        )
        result = self.fetch_resume_import(connection, session_id, import_id)
        assert result is not None
        return result

    def fetch_resume_import(
        self, connection: sqlite3.Connection, session_id: str, import_id: str
    ) -> dict[str, Any] | None:
        row = connection.execute(
            "SELECT * FROM resume_imports WHERE id = ? AND session_id = ?", (import_id, session_id)
        ).fetchone()
        if row is None:
            return None
        result = dict(row)
        result["document"] = json.loads(result.pop("document_json")) if result.get("document_json") else None
        result["field_reviews"] = json.loads(result.pop("field_reviews_json"))
        return result

    def find_resume_import_by_sha256(
        self, connection: sqlite3.Connection, session_id: str, sha256: str
    ) -> dict[str, Any] | None:
        row = connection.execute(
            "SELECT id FROM resume_imports WHERE session_id = ? AND sha256 = ?",
            (session_id, sha256),
        ).fetchone()
        return self.fetch_resume_import(connection, session_id, row["id"]) if row else None


    def update_resume_import_status(
        self, connection: sqlite3.Connection, session_id: str, import_id: str, status: str
    ) -> dict[str, Any]:
        connection.execute(
            "UPDATE resume_imports SET status = ?, updated_at = ? WHERE id = ? AND session_id = ?",
            (status, utc_now(), import_id, session_id),
        )
        result = self.fetch_resume_import(connection, session_id, import_id)
        if result is None:
            raise KeyError(import_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=attach_gap_report_staleness(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