from __future__ import annotations

from sqlalchemy import text

from app.modules.rag.persistence.schema_statements import legacy_chunk_columns


class RagSchemaMigrator:
    def migrate(self, conn) -> None:
        self._ensure_core_columns(conn)
        self._convert_metadata_to_jsonb(conn)
        self._migrate_legacy_chunk_metadata(conn)
        self._migrate_legacy_cache(conn)
        self._drop_legacy_chunk_columns(conn)

    def _ensure_core_columns(self, conn) -> None:
        for statement in (
            "ALTER TABLE rag_chunks ADD COLUMN IF NOT EXISTS layer VARCHAR(64) NULL",
            "ALTER TABLE rag_chunks ADD COLUMN IF NOT EXISTS lang VARCHAR(32) NULL",
            "ALTER TABLE rag_chunks ADD COLUMN IF NOT EXISTS repo_id VARCHAR(512) NULL",
            "ALTER TABLE rag_chunks ADD COLUMN IF NOT EXISTS commit_sha VARCHAR(128) NULL",
            "ALTER TABLE rag_chunks ADD COLUMN IF NOT EXISTS title TEXT NULL",
            "ALTER TABLE rag_chunks ADD COLUMN IF NOT EXISTS metadata_json JSONB NOT NULL DEFAULT '{}'::jsonb",
            "ALTER TABLE rag_chunks ADD COLUMN IF NOT EXISTS span_start INTEGER NULL",
            "ALTER TABLE rag_chunks ADD COLUMN IF NOT EXISTS span_end INTEGER NULL",
        ):
            conn.execute(text(statement))

    def _convert_metadata_to_jsonb(self, conn) -> None:
        conn.execute(
            text(
                """
                ALTER TABLE rag_chunks
                ALTER COLUMN metadata_json TYPE JSONB
                USING CASE
                    WHEN metadata_json IS NULL THEN '{}'::jsonb
                    WHEN trim(metadata_json::text) = '' THEN '{}'::jsonb
                    ELSE metadata_json::jsonb
                END
                """
            )
        )

    def _migrate_legacy_chunk_metadata(self, conn) -> None:
        pairs = self._legacy_chunk_pairs(conn)
        if not pairs:
            return
        entries = ",\n                        ".join(f"'{key}', {column}" for key, column in pairs)
        conn.execute(
            text(
                f"""
                UPDATE rag_chunks
                SET metadata_json = COALESCE(metadata_json, '{{}}'::jsonb) || jsonb_strip_nulls(
                    jsonb_build_object(
                        {entries}
                    )
                )
                """
            )
        )

    def _migrate_legacy_cache(self, conn) -> None:
        row = conn.execute(text("SELECT to_regclass('public.rag_chunk_cache') AS table_name")).mappings().first()
        if not row or not row.get("table_name"):
            return
        conn.execute(
            text(
                """
                INSERT INTO rag_document_cache (
                    repo_id, blob_sha, chunk_index, path, layer, lang, title, content, embedding,
                    metadata_json, links_json, span_start, span_end, commit_sha, created_at, updated_at
                )
                SELECT
                    repo_id,
                    blob_sha,
                    chunk_index,
                    COALESCE(path, ''),
                    layer,
                    lang,
                    title,
                    content,
                    embedding,
                    CASE
                        WHEN metadata_json IS NULL THEN '{}'::jsonb
                        WHEN trim(metadata_json) = '' THEN '{}'::jsonb
                        ELSE metadata_json::jsonb
                    END,
                    CASE
                        WHEN links_json IS NULL THEN '[]'::jsonb
                        WHEN trim(links_json) = '' THEN '[]'::jsonb
                        ELSE links_json::jsonb
                    END,
                    span_start,
                    span_end,
                    commit_sha,
                    created_at,
                    updated_at
                FROM rag_chunk_cache
                ON CONFLICT (repo_id, blob_sha, chunk_index) DO NOTHING
                """
            )
        )

    def _drop_legacy_chunk_columns(self, conn) -> None:
        for column in legacy_chunk_columns():
            conn.execute(text(f"ALTER TABLE rag_chunks DROP COLUMN IF EXISTS {column}"))

    def _legacy_chunk_pairs(self, conn) -> list[tuple[str, str]]:
        desired = (
            ("artifact_type", "artifact_type"),
            ("section", "section"),
            ("doc_id", "doc_id"),
            ("doc_version", "doc_version"),
            ("owner", "owner"),
            ("system_component", "system_component"),
            ("last_modified", "last_modified"),
            ("staleness_score", "staleness_score"),
            ("symbol_id", "symbol_id"),
            ("qname", "qname"),
            ("kind", "kind"),
            ("framework", "framework"),
            ("entrypoint_type", "entrypoint_type"),
            ("module_id", "module_id"),
            ("section_path", "section_path"),
            ("doc_kind", "doc_kind"),
        )
        rows = conn.execute(
            text(
                """
                SELECT column_name
                FROM information_schema.columns
                WHERE table_schema = 'public' AND table_name = 'rag_chunks'
                """
            )
        ).fetchall()
        existing = {str(row[0]) for row in rows}
        return [(key, column) for key, column in desired if column in existing]
