"""Implement infrastructure support for postgres document processing history."""

from __future__ import annotations

import json
from contextlib import contextmanager
from datetime import UTC, datetime
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 DocumentArtifact, DocumentProcessingRun, DocumentProcessingStore, DocumentStatusEvent
# Keep SQL mapping local to this adapter so the domain stays storage-agnostic.

_CREATE_RUNS_TABLE = """
CREATE TABLE IF NOT EXISTS document_processing_runs (
    run_id                 VARCHAR(128) PRIMARY KEY,
    doc_id                 VARCHAR(128) NOT NULL,
    trigger_type           VARCHAR(32) NOT NULL,
    run_status             VARCHAR(32) NOT NULL DEFAULT 'running',
    parser_backend         VARCHAR(128) NOT NULL DEFAULT '',
    chunk_backend          VARCHAR(128) NOT NULL DEFAULT '',
    embedding_model        VARCHAR(256) NOT NULL DEFAULT '',
    index_name             VARCHAR(128) NOT NULL DEFAULT '',
    started_at             TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    stored_at              TIMESTAMPTZ,
    parsed_at              TIMESTAMPTZ,
    indexed_at             TIMESTAMPTZ,
    finished_at            TIMESTAMPTZ,
    layout_count           INTEGER NOT NULL DEFAULT 0,
    structure_node_count   INTEGER NOT NULL DEFAULT 0,
    semantic_block_count   INTEGER NOT NULL DEFAULT 0,
    vector_chunk_count     INTEGER NOT NULL DEFAULT 0,
    chunk_count            INTEGER NOT NULL DEFAULT 0,
    failure_stage          VARCHAR(64) NOT NULL DEFAULT '',
    error_message          TEXT NOT NULL DEFAULT '',
    metadata               JSONB NOT NULL DEFAULT '{}',
    created_at             TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at             TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    CONSTRAINT fk_dpr_doc FOREIGN KEY (doc_id) REFERENCES documents(doc_id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_document_processing_runs_doc_id ON document_processing_runs(doc_id, started_at DESC);
"""

_CREATE_EVENTS_TABLE = """
CREATE TABLE IF NOT EXISTS document_status_history (
    event_id               VARCHAR(128) PRIMARY KEY,
    doc_id                 VARCHAR(128) NOT NULL,
    run_id                 VARCHAR(128) NOT NULL,
    from_status            VARCHAR(32) NOT NULL DEFAULT '',
    to_status              VARCHAR(32) NOT NULL,
    stage                  VARCHAR(64) NOT NULL DEFAULT '',
    message                TEXT NOT NULL DEFAULT '',
    metadata               JSONB NOT NULL DEFAULT '{}',
    occurred_at            TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    CONSTRAINT fk_dsh_doc FOREIGN KEY (doc_id) REFERENCES documents(doc_id) ON DELETE CASCADE,
    CONSTRAINT fk_dsh_run FOREIGN KEY (run_id) REFERENCES document_processing_runs(run_id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_document_status_history_doc_id ON document_status_history(doc_id, occurred_at ASC);
CREATE INDEX IF NOT EXISTS idx_document_status_history_run_id ON document_status_history(run_id, occurred_at ASC);
"""

_CREATE_ARTIFACTS_TABLE = """
CREATE TABLE IF NOT EXISTS document_artifacts (
    artifact_id            VARCHAR(128) PRIMARY KEY,
    doc_id                 VARCHAR(128) NOT NULL,
    run_id                 VARCHAR(128) NOT NULL,
    artifact_type          VARCHAR(64) NOT NULL,
    object_name            VARCHAR(1024) NOT NULL,
    content_type           VARCHAR(128) NOT NULL DEFAULT '',
    byte_size              BIGINT NOT NULL DEFAULT 0,
    checksum               VARCHAR(256) NOT NULL DEFAULT '',
    metadata               JSONB NOT NULL DEFAULT '{}',
    created_at             TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    CONSTRAINT fk_da_doc FOREIGN KEY (doc_id) REFERENCES documents(doc_id) ON DELETE CASCADE,
    CONSTRAINT fk_da_run FOREIGN KEY (run_id) REFERENCES document_processing_runs(run_id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_document_artifacts_doc_id ON document_artifacts(doc_id, created_at ASC);
CREATE INDEX IF NOT EXISTS idx_document_artifacts_run_id ON document_artifacts(run_id, created_at ASC);
"""


class PostgresDocumentProcessingStore(DocumentProcessingStore):
    """Persist processing history in PostgreSQL using handwritten SQL."""

    def __init__(self) -> None:
        """Initialize the store and ensure the required tables exist."""
        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:
        """Create processing history tables and indexes if they are missing."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(_CREATE_RUNS_TABLE)
                cur.execute(_CREATE_EVENTS_TABLE)
                cur.execute(_CREATE_ARTIFACTS_TABLE)
            conn.commit()

    @contextmanager
    def _conn(self):
        """Borrow one connection from the pool and return it afterwards."""
        conn = self._pool.getconn()
        try:
            yield conn
        finally:
            self._pool.putconn(conn)

    def _normalize_metadata(self, value: Any) -> dict[str, Any]:
        """Return a JSON-object payload regardless of the row representation."""
        if isinstance(value, dict):
            return value
        if not value:
            return {}
        return json.loads(value)

    def _row_to_run(self, row: dict[str, Any]) -> DocumentProcessingRun:
        """Map one run row into the domain dataclass."""
        return DocumentProcessingRun(
            run_id=row["run_id"],
            doc_id=row["doc_id"],
            trigger_type=row["trigger_type"],
            run_status=row["run_status"],
            parser_backend=row["parser_backend"],
            chunk_backend=row["chunk_backend"],
            embedding_model=row["embedding_model"],
            index_name=row["index_name"],
            started_at=row["started_at"],
            stored_at=row["stored_at"],
            parsed_at=row["parsed_at"],
            indexed_at=row["indexed_at"],
            finished_at=row["finished_at"],
            layout_count=row["layout_count"],
            structure_node_count=row["structure_node_count"],
            semantic_block_count=row["semantic_block_count"],
            vector_chunk_count=row["vector_chunk_count"],
            chunk_count=row["chunk_count"],
            failure_stage=row["failure_stage"],
            error_message=row["error_message"],
            metadata=self._normalize_metadata(row["metadata"]),
        )

    def _row_to_event(self, row: dict[str, Any]) -> DocumentStatusEvent:
        """Map one event row into the domain dataclass."""
        return DocumentStatusEvent(
            event_id=row["event_id"],
            doc_id=row["doc_id"],
            run_id=row["run_id"],
            from_status=row["from_status"],
            to_status=row["to_status"],
            stage=row["stage"],
            message=row["message"],
            metadata=self._normalize_metadata(row["metadata"]),
            occurred_at=row["occurred_at"],
        )

    def _row_to_artifact(self, row: dict[str, Any]) -> DocumentArtifact:
        """Map one artifact row into the domain dataclass."""
        return DocumentArtifact(
            artifact_id=row["artifact_id"],
            doc_id=row["doc_id"],
            run_id=row["run_id"],
            artifact_type=row["artifact_type"],
            object_name=row["object_name"],
            content_type=row["content_type"],
            byte_size=row["byte_size"],
            checksum=row["checksum"],
            metadata=self._normalize_metadata(row["metadata"]),
            created_at=row["created_at"],
        )

    def _update_run(
        self,
        run_id: str,
        *,
        assignments: dict[str, Any],
        metadata: dict | None = None,
    ) -> DocumentProcessingRun | None:
        """Update one run row and return the latest stored state."""
        set_clauses = []
        params: dict[str, Any] = {"run_id": run_id, "updated_at": datetime.now(UTC)}
        for key, value in assignments.items():
            set_clauses.append(f"{key} = %({key})s")
            params[key] = value
        set_clauses.append("updated_at = %(updated_at)s")
        if metadata is not None:
            set_clauses.append("metadata = COALESCE(metadata, '{}'::jsonb) || %(metadata)s::jsonb")
            params["metadata"] = json.dumps(metadata, ensure_ascii=False)
        sql = f"""
            UPDATE document_processing_runs
            SET {", ".join(set_clauses)}
            WHERE run_id = %(run_id)s
            RETURNING *
        """
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, params)
                row = cur.fetchone()
            conn.commit()
        return self._row_to_run(dict(row)) if row else None

    def create_run(self, run: DocumentProcessingRun) -> DocumentProcessingRun:
        """Create a new processing run record."""
        sql = """
            INSERT INTO document_processing_runs
                (run_id, doc_id, trigger_type, run_status, parser_backend, chunk_backend,
                 embedding_model, index_name, started_at, stored_at, parsed_at, indexed_at,
                 finished_at, layout_count, structure_node_count, semantic_block_count,
                 vector_chunk_count, chunk_count, failure_stage, error_message, metadata)
            VALUES
                (%(run_id)s, %(doc_id)s, %(trigger_type)s, %(run_status)s, %(parser_backend)s,
                 %(chunk_backend)s, %(embedding_model)s, %(index_name)s, %(started_at)s,
                 %(stored_at)s, %(parsed_at)s, %(indexed_at)s, %(finished_at)s, %(layout_count)s,
                 %(structure_node_count)s, %(semantic_block_count)s, %(vector_chunk_count)s,
                 %(chunk_count)s, %(failure_stage)s, %(error_message)s, %(metadata)s)
        """
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    sql,
                    {
                        "run_id": run.run_id,
                        "doc_id": run.doc_id,
                        "trigger_type": run.trigger_type,
                        "run_status": run.run_status,
                        "parser_backend": run.parser_backend,
                        "chunk_backend": run.chunk_backend,
                        "embedding_model": run.embedding_model,
                        "index_name": run.index_name,
                        "started_at": run.started_at,
                        "stored_at": run.stored_at,
                        "parsed_at": run.parsed_at,
                        "indexed_at": run.indexed_at,
                        "finished_at": run.finished_at,
                        "layout_count": run.layout_count,
                        "structure_node_count": run.structure_node_count,
                        "semantic_block_count": run.semantic_block_count,
                        "vector_chunk_count": run.vector_chunk_count,
                        "chunk_count": run.chunk_count,
                        "failure_stage": run.failure_stage,
                        "error_message": run.error_message,
                        "metadata": json.dumps(run.metadata, ensure_ascii=False),
                    },
                )
            conn.commit()
        return run

    def mark_run_stored(
        self,
        run_id: str,
        *,
        stored_at: datetime | None = None,
        metadata: dict | None = None,
    ) -> DocumentProcessingRun | None:
        """Mark a run as having persisted its source file."""
        return self._update_run(
            run_id,
            assignments={"stored_at": stored_at or datetime.now(UTC)},
            metadata=metadata,
        )

    def mark_run_parsed(
        self,
        run_id: str,
        *,
        parser_backend: str,
        layout_count: int,
        structure_node_count: int,
        semantic_block_count: int,
        vector_chunk_count: int,
        parsed_at: datetime | None = None,
        metadata: dict | None = None,
    ) -> DocumentProcessingRun | None:
        """Record parse completion metrics for a run."""
        return self._update_run(
            run_id,
            assignments={
                "parser_backend": parser_backend,
                "parsed_at": parsed_at or datetime.now(UTC),
                "layout_count": layout_count,
                "structure_node_count": structure_node_count,
                "semantic_block_count": semantic_block_count,
                "vector_chunk_count": vector_chunk_count,
            },
            metadata=metadata,
        )

    def mark_run_indexed(
        self,
        run_id: str,
        *,
        chunk_count: int,
        index_name: str,
        indexed_at: datetime | None = None,
        finished_at: datetime | None = None,
        metadata: dict | None = None,
    ) -> DocumentProcessingRun | None:
        """Mark a run as successfully indexed."""
        now = datetime.now(UTC)
        return self._update_run(
            run_id,
            assignments={
                "run_status": "succeeded",
                "chunk_count": chunk_count,
                "index_name": index_name,
                "indexed_at": indexed_at or now,
                "finished_at": finished_at or now,
            },
            metadata=metadata,
        )

    def mark_run_failed(
        self,
        run_id: str,
        *,
        failure_stage: str,
        error_message: str,
        finished_at: datetime | None = None,
        metadata: dict | None = None,
    ) -> DocumentProcessingRun | None:
        """Mark a run as failed and persist the terminal error details."""
        return self._update_run(
            run_id,
            assignments={
                "run_status": "failed",
                "failure_stage": failure_stage,
                "error_message": error_message,
                "finished_at": finished_at or datetime.now(UTC),
            },
            metadata=metadata,
        )

    def append_status_event(self, event: DocumentStatusEvent) -> DocumentStatusEvent:
        """Append a document status event."""
        sql = """
            INSERT INTO document_status_history
                (event_id, doc_id, run_id, from_status, to_status, stage, message, metadata, occurred_at)
            VALUES
                (%(event_id)s, %(doc_id)s, %(run_id)s, %(from_status)s, %(to_status)s,
                 %(stage)s, %(message)s, %(metadata)s, %(occurred_at)s)
        """
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    sql,
                    {
                        "event_id": event.event_id,
                        "doc_id": event.doc_id,
                        "run_id": event.run_id,
                        "from_status": event.from_status,
                        "to_status": event.to_status,
                        "stage": event.stage,
                        "message": event.message,
                        "metadata": json.dumps(event.metadata, ensure_ascii=False),
                        "occurred_at": event.occurred_at,
                    },
                )
            conn.commit()
        return event

    def replace_artifacts_for_run(self, run_id: str, artifacts: list[DocumentArtifact]) -> list[DocumentArtifact]:
        """Replace all artifact references for one run using a delete-then-insert strategy."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM document_artifacts WHERE run_id = %s", (run_id,))
                if artifacts:
                    psycopg2.extras.execute_values(
                        cur,
                        """
                        INSERT INTO document_artifacts
                            (artifact_id, doc_id, run_id, artifact_type, object_name,
                             content_type, byte_size, checksum, metadata, created_at)
                        VALUES %s
                        """,
                        [
                            (
                                artifact.artifact_id,
                                artifact.doc_id,
                                artifact.run_id,
                                artifact.artifact_type,
                                artifact.object_name,
                                artifact.content_type,
                                artifact.byte_size,
                                artifact.checksum,
                                json.dumps(artifact.metadata, ensure_ascii=False),
                                artifact.created_at,
                            )
                            for artifact in artifacts
                        ],
                    )
            conn.commit()
        return artifacts

    def delete_by_document(self, doc_id: str) -> None:
        """Delete all processing rows for a document explicitly."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM document_status_history WHERE doc_id = %s", (doc_id,))
                cur.execute("DELETE FROM document_artifacts WHERE doc_id = %s", (doc_id,))
                cur.execute("DELETE FROM document_processing_runs WHERE doc_id = %s", (doc_id,))
            conn.commit()

    def list_runs_by_document(self, doc_id: str) -> list[DocumentProcessingRun]:
        """List processing runs for a document in chronological order."""
        sql = "SELECT * FROM document_processing_runs WHERE doc_id = %s ORDER BY started_at ASC"
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, (doc_id,))
                rows = cur.fetchall()
        return [self._row_to_run(dict(row)) for row in rows]

    def get_run(self, run_id: str) -> DocumentProcessingRun | None:
        """Return one processing run by identifier."""
        sql = "SELECT * FROM document_processing_runs WHERE run_id = %s"
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, (run_id,))
                row = cur.fetchone()
        return self._row_to_run(dict(row)) if row else None

    def list_status_events_by_document(self, doc_id: str) -> list[DocumentStatusEvent]:
        """List all status events for a document."""
        sql = "SELECT * FROM document_status_history WHERE doc_id = %s ORDER BY occurred_at ASC"
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, (doc_id,))
                rows = cur.fetchall()
        return [self._row_to_event(dict(row)) for row in rows]

    def list_status_events_by_run(self, run_id: str) -> list[DocumentStatusEvent]:
        """List all status events for a run."""
        sql = "SELECT * FROM document_status_history WHERE run_id = %s ORDER BY occurred_at ASC"
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, (run_id,))
                rows = cur.fetchall()
        return [self._row_to_event(dict(row)) for row in rows]

    def list_artifacts_by_document(self, doc_id: str) -> list[DocumentArtifact]:
        """List all artifact references for a document."""
        sql = "SELECT * FROM document_artifacts WHERE doc_id = %s ORDER BY created_at ASC"
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, (doc_id,))
                rows = cur.fetchall()
        return [self._row_to_artifact(dict(row)) for row in rows]

    def list_artifacts_by_run(self, run_id: str) -> list[DocumentArtifact]:
        """List all artifact references for a run."""
        sql = "SELECT * FROM document_artifacts WHERE run_id = %s ORDER BY created_at ASC"
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, (run_id,))
                rows = cur.fetchall()
        return [self._row_to_artifact(dict(row)) for row in rows]
