"""Implement infrastructure support for postgres parse artifact store."""

from __future__ import annotations

import json
from contextlib import contextmanager
from typing import Any

import psycopg2
import psycopg2.extras
from psycopg2.pool import ThreadedConnectionPool

from app.config.settings import settings
from app.domain.documents import ParseArtifactStore

_CREATE_STRUCTURE_NODES = """
CREATE TABLE IF NOT EXISTS structure_nodes (
    id        SERIAL PRIMARY KEY,
    doc_id    VARCHAR(128) NOT NULL,
    unique_id VARCHAR(128),
    page      INTEGER      NOT NULL DEFAULT 0,
    idx       INTEGER      NOT NULL DEFAULT 0,
    level     INTEGER      NOT NULL DEFAULT 0,
    title     TEXT         NOT NULL DEFAULT '',
    type      VARCHAR(64),
    sub_type  VARCHAR(64),
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    CONSTRAINT fk_sn_doc FOREIGN KEY (doc_id) REFERENCES documents(doc_id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_structure_nodes_doc_id ON structure_nodes(doc_id);
"""

_CREATE_SEMANTIC_BLOCKS = """
CREATE TABLE IF NOT EXISTS semantic_blocks (
    id           SERIAL PRIMARY KEY,
    doc_id       VARCHAR(128) NOT NULL,
    semantic_id  VARCHAR(128) NOT NULL,
    block_type   VARCHAR(64)  NOT NULL DEFAULT '',
    page_start   INTEGER      NOT NULL DEFAULT 0,
    page_end     INTEGER      NOT NULL DEFAULT 0,
    section_path JSONB        NOT NULL DEFAULT '[]',
    section_level INTEGER     NOT NULL DEFAULT 0,
    section_title VARCHAR(512) NOT NULL DEFAULT '',
    source_ids   JSONB        NOT NULL DEFAULT '[]',
    text         TEXT         NOT NULL DEFAULT '',
    created_at   TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
    CONSTRAINT fk_sb_doc FOREIGN KEY (doc_id) REFERENCES documents(doc_id) ON DELETE CASCADE,
    CONSTRAINT uq_semantic_blocks UNIQUE (doc_id, semantic_id)
);
CREATE INDEX IF NOT EXISTS idx_semantic_blocks_doc_id ON semantic_blocks(doc_id);
"""


class PostgresParseArtifactStore(ParseArtifactStore):
    """ParseArtifactStore implementation backed by PostgreSQL.

    Requires the `documents` table to exist first (created by PostgresDocumentRepository).
    """

    def __init__(self) -> None:
        self._pool = ThreadedConnectionPool(
            minconn=1,
            maxconn=5,
            host=settings.postgres_host,
            port=settings.postgres_port,
            user=settings.postgres_user,
            password=settings.postgres_password,
            dbname=settings.postgres_db,
        )
        self._ensure_schema()

    def _ensure_schema(self) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(_CREATE_STRUCTURE_NODES)
                cur.execute(_CREATE_SEMANTIC_BLOCKS)
            conn.commit()

    @contextmanager
    def _conn(self):
        conn = self._pool.getconn()
        try:
            yield conn
        finally:
            self._pool.putconn(conn)

    # ------------------------------------------------------------------
    # ParseArtifactStore interface
    # ------------------------------------------------------------------

    def save(
        self,
        doc_id: str,
        structure_nodes: list[dict],
        semantic_blocks: list[dict],
    ) -> None:
        """Persist structure nodes and semantic blocks, replacing any existing records."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                # Delete existing records first to keep save idempotent.
                cur.execute("DELETE FROM structure_nodes WHERE doc_id = %s", (doc_id,))
                cur.execute("DELETE FROM semantic_blocks WHERE doc_id = %s", (doc_id,))

                if structure_nodes:
                    psycopg2.extras.execute_values(
                        cur,
                        """
                        INSERT INTO structure_nodes
                            (doc_id, unique_id, page, idx, level, title, type, sub_type)
                        VALUES %s
                        """,
                        [
                            (
                                doc_id,
                                node.get("unique_id"),
                                int(node.get("page", 0) or 0),
                                int(node.get("index", 0) or 0),
                                int(node.get("level", 0) or 0),
                                str(node.get("title", "")),
                                node.get("type"),
                                node.get("sub_type"),
                            )
                            for node in structure_nodes
                        ],
                    )

                if semantic_blocks:
                    psycopg2.extras.execute_values(
                        cur,
                        """
                        INSERT INTO semantic_blocks
                            (doc_id, semantic_id, block_type, page_start, page_end,
                             section_path, section_level, section_title, source_ids, text)
                        VALUES %s
                        """,
                        [
                            (
                                doc_id,
                                block.get("semantic_id", ""),
                                block.get("block_type", ""),
                                int(block.get("page_start", 0) or 0),
                                int(block.get("page_end", 0) or 0),
                                json.dumps(block.get("section_path", []), ensure_ascii=False),
                                int(block.get("section_level", 0) or 0),
                                str(block.get("section_title", "")),
                                json.dumps(block.get("source_ids", []), ensure_ascii=False),
                                str(block.get("text", "")),
                            )
                            for block in semantic_blocks
                        ],
                    )
            conn.commit()

    def delete(self, doc_id: str) -> None:
        """Remove all parse artifacts for a document (ON DELETE CASCADE handles child rows)."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM structure_nodes WHERE doc_id = %s", (doc_id,))
                cur.execute("DELETE FROM semantic_blocks WHERE doc_id = %s", (doc_id,))
            conn.commit()

    def get_semantic_blocks(self, doc_id: str) -> list[dict[str, Any]]:
        """Return all semantic blocks for a document ordered by id."""
        sql = """
            SELECT semantic_id, block_type, page_start, page_end,
                   section_path, section_level, section_title, source_ids, text
            FROM semantic_blocks
            WHERE doc_id = %s
            ORDER BY id
        """
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, (doc_id,))
                rows = cur.fetchall()
        results = []
        for row in rows:
            item = dict(row)
            for key in ("section_path", "source_ids"):
                if isinstance(item[key], str):
                    item[key] = json.loads(item[key])
            results.append(item)
        return results

    def get_structure_nodes(self, doc_id: str) -> list[dict[str, Any]]:
        """Return all structure nodes for a document ordered by idx."""
        sql = """
            SELECT unique_id, page, idx, level, title, type, sub_type
            FROM structure_nodes
            WHERE doc_id = %s
            ORDER BY idx
        """
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, (doc_id,))
                rows = cur.fetchall()
        return [dict(r) for r in rows]
