"""add retrieval service v1

Revision ID: c3d4e5f6a7b8
Revises: b2c3d4e5f6a7
Create Date: 2026-04-19 11:19:17.905181

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = "c3d4e5f6a7b8"
down_revision: Union[str, Sequence[str], None] = "b2c3d4e5f6a7"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""
    op.execute(
        """
        CREATE UNIQUE INDEX IF NOT EXISTS uq_jobs_user_active_document
        ON jobs USING btree (user_id, ((job_metadata ->> 'document_id')))
        WHERE (
            status IN ('waiting-file', 'pending', 'running', 'converting')
            AND (job_metadata ->> 'document_id') IS NOT NULL
        )
        """
    )
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "documents",
        sa.Column("document_id", sa.String(length=36), nullable=False),
        sa.Column("user_id", sa.Text(), nullable=False),
        sa.Column("namespace", sa.String(length=255), nullable=False),
        sa.Column("status", sa.String(length=32), nullable=False),
        sa.Column("current_job_result_id", sa.String(length=36), nullable=True),
        sa.Column("source_file_name", sa.Text(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.Column("archived_at", sa.DateTime(), nullable=True),
        sa.ForeignKeyConstraint(
            ["current_job_result_id"], ["job_results.id"], ondelete="SET NULL"
        ),
        sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="RESTRICT"),
        sa.PrimaryKeyConstraint("document_id"),
    )
    op.create_index(
        "idx_documents_current_job_result",
        "documents",
        ["current_job_result_id"],
        unique=False,
    )
    op.create_index(
        "idx_documents_user_namespace_status",
        "documents",
        ["user_id", "namespace", "status"],
        unique=False,
    )
    op.create_table(
        "document_sections",
        sa.Column("section_id", sa.String(length=36), nullable=False),
        sa.Column("user_id", sa.Text(), nullable=False),
        sa.Column("namespace", sa.String(length=255), nullable=False),
        sa.Column("document_id", sa.String(length=36), nullable=False),
        sa.Column("job_result_id", sa.String(length=36), nullable=False),
        sa.Column("parent_section_id", sa.String(length=36), nullable=True),
        sa.Column("section_path", sa.Text(), nullable=False),
        sa.Column("section_title", sa.Text(), nullable=True),
        sa.Column("section_level", sa.Integer(), nullable=False),
        sa.Column("summary", sa.Text(), nullable=True),
        sa.Column("section_metadata", sa.JSON(), nullable=True),
        sa.Column("sort_order", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["document_id"], ["documents.document_id"], ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["job_result_id"], ["job_results.id"], ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["parent_section_id"], ["document_sections.section_id"], ondelete="SET NULL"
        ),
        sa.PrimaryKeyConstraint("section_id"),
        sa.UniqueConstraint(
            "document_id",
            "job_result_id",
            "section_path",
            name="uq_document_sections_revision_path",
        ),
    )
    op.create_index(
        "idx_document_sections_doc_revision",
        "document_sections",
        ["document_id", "job_result_id"],
        unique=False,
    )
    op.create_index(
        "idx_document_sections_scope",
        "document_sections",
        ["user_id", "namespace"],
        unique=False,
    )
    op.create_table(
        "graph_nodes",
        sa.Column("node_id", sa.String(length=128), nullable=False),
        sa.Column("user_id", sa.Text(), nullable=False),
        sa.Column("namespace", sa.String(length=255), nullable=False),
        sa.Column("node_kind", sa.String(length=32), nullable=False),
        sa.Column("owner_document_id", sa.String(length=36), nullable=False),
        sa.Column("job_result_id", sa.String(length=36), nullable=False),
        sa.Column("ref_document_id", sa.String(length=36), nullable=True),
        sa.Column("ref_section_id", sa.String(length=36), nullable=True),
        sa.Column("properties", sa.JSON(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["job_result_id"], ["job_results.id"], ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["owner_document_id"], ["documents.document_id"], ondelete="CASCADE"
        ),
        sa.PrimaryKeyConstraint("node_id"),
    )
    op.create_index(
        "idx_graph_nodes_owner_revision",
        "graph_nodes",
        ["owner_document_id", "job_result_id"],
        unique=False,
    )
    op.create_index(
        "idx_graph_nodes_ref_document", "graph_nodes", ["ref_document_id"], unique=False
    )
    op.create_index(
        "idx_graph_nodes_ref_section", "graph_nodes", ["ref_section_id"], unique=False
    )
    op.create_index(
        "idx_graph_nodes_scope",
        "graph_nodes",
        ["user_id", "namespace", "node_kind"],
        unique=False,
    )
    op.create_table(
        "retrieval_hit_stats",
        sa.Column("id", sa.String(length=36), nullable=False),
        sa.Column("user_id", sa.Text(), nullable=False),
        sa.Column("namespace", sa.String(length=255), nullable=False),
        sa.Column("hit_kind", sa.String(length=32), nullable=False),
        sa.Column("document_id", sa.String(length=36), nullable=False),
        sa.Column("chunk_id", sa.String(length=64), nullable=True),
        sa.Column("hit_count", sa.Integer(), nullable=False),
        sa.Column("last_hit_at", sa.DateTime(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["document_id"], ["documents.document_id"], ondelete="CASCADE"
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        "idx_retrieval_hit_stats_chunk",
        "retrieval_hit_stats",
        ["chunk_id"],
        unique=False,
    )
    op.create_index(
        "idx_retrieval_hit_stats_document",
        "retrieval_hit_stats",
        ["document_id"],
        unique=False,
    )
    op.create_index(
        "idx_retrieval_hit_stats_scope_kind",
        "retrieval_hit_stats",
        ["user_id", "namespace", "hit_kind"],
        unique=False,
    )
    op.create_index(
        "uq_retrieval_hit_stats_document_key",
        "retrieval_hit_stats",
        ["user_id", "namespace", "hit_kind", "document_id"],
        unique=True,
        postgresql_where=sa.text("chunk_id IS NULL"),
    )
    op.create_index(
        "uq_retrieval_hit_stats_chunk_key",
        "retrieval_hit_stats",
        ["user_id", "namespace", "hit_kind", "document_id", "chunk_id"],
        unique=True,
        postgresql_where=sa.text("chunk_id IS NOT NULL"),
    )
    op.create_table(
        "document_chunks",
        sa.Column("id", sa.String(length=36), nullable=False),
        sa.Column("chunk_id", sa.String(length=64), nullable=False),
        sa.Column("user_id", sa.Text(), nullable=False),
        sa.Column("namespace", sa.String(length=255), nullable=False),
        sa.Column("document_id", sa.String(length=36), nullable=False),
        sa.Column("job_result_id", sa.String(length=36), nullable=False),
        sa.Column("section_id", sa.String(length=36), nullable=True),
        sa.Column("chunk_type", sa.String(length=64), nullable=False),
        sa.Column("content", sa.Text(), nullable=True),
        sa.Column("content_lexical_text", sa.Text(), nullable=True),
        sa.Column("path_lexical_text", sa.Text(), nullable=True),
        sa.Column("source_chunk_path", sa.Text(), nullable=True),
        sa.Column("file_path", sa.Text(), nullable=True),
        sa.Column("chunk_metadata", sa.JSON(), nullable=True),
        sa.Column("sort_order", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["document_id"], ["documents.document_id"], ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["job_result_id"], ["job_results.id"], ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["section_id"], ["document_sections.section_id"], ondelete="SET NULL"
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "document_id",
            "job_result_id",
            "source_chunk_path",
            name="uq_document_chunks_revision_path",
        ),
    )
    op.create_index(
        "idx_document_chunks_chunk_id", "document_chunks", ["chunk_id"], unique=False
    )
    op.create_index(
        "idx_document_chunks_doc_revision",
        "document_chunks",
        ["document_id", "job_result_id"],
        unique=False,
    )
    op.create_index(
        "idx_document_chunks_scope",
        "document_chunks",
        ["user_id", "namespace"],
        unique=False,
    )
    op.create_index(
        "idx_document_chunks_section", "document_chunks", ["section_id"], unique=False
    )
    op.create_table(
        "graph_edges",
        sa.Column("edge_id", sa.String(length=160), nullable=False),
        sa.Column("user_id", sa.Text(), nullable=False),
        sa.Column("namespace", sa.String(length=255), nullable=False),
        sa.Column("edge_kind", sa.String(length=32), nullable=False),
        sa.Column("source_node_id", sa.String(length=128), nullable=False),
        sa.Column("target_node_id", sa.String(length=128), nullable=False),
        sa.Column("owner_document_id", sa.String(length=36), nullable=False),
        sa.Column("job_result_id", sa.String(length=36), nullable=False),
        sa.Column("is_directed", sa.Boolean(), nullable=False),
        sa.Column("weight", sa.Float(), nullable=True),
        sa.Column("properties", sa.JSON(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["job_result_id"], ["job_results.id"], ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["owner_document_id"], ["documents.document_id"], ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["source_node_id"], ["graph_nodes.node_id"], ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(
            ["target_node_id"], ["graph_nodes.node_id"], ondelete="CASCADE"
        ),
        sa.PrimaryKeyConstraint("edge_id"),
    )
    op.create_index(
        "idx_graph_edges_owner_revision",
        "graph_edges",
        ["owner_document_id", "job_result_id"],
        unique=False,
    )
    op.create_index(
        "idx_graph_edges_scope",
        "graph_edges",
        ["user_id", "namespace", "edge_kind"],
        unique=False,
    )
    op.create_index(
        "idx_graph_edges_source", "graph_edges", ["source_node_id"], unique=False
    )
    op.create_index(
        "idx_graph_edges_target", "graph_edges", ["target_node_id"], unique=False
    )
    op.add_column(
        "job_results", sa.Column("document_id", sa.String(length=36), nullable=True)
    )
    op.create_index(
        op.f("ix_job_results_document_id"), "job_results", ["document_id"], unique=False
    )
    op.create_foreign_key(
        "fk_job_results_document_id",
        "job_results",
        "documents",
        ["document_id"],
        ["document_id"],
        ondelete="SET NULL",
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint("fk_job_results_document_id", "job_results", type_="foreignkey")
    op.drop_index(op.f("ix_job_results_document_id"), table_name="job_results")
    op.drop_column("job_results", "document_id")
    op.drop_index("idx_graph_edges_target", table_name="graph_edges")
    op.drop_index("idx_graph_edges_source", table_name="graph_edges")
    op.drop_index("idx_graph_edges_scope", table_name="graph_edges")
    op.drop_index("idx_graph_edges_owner_revision", table_name="graph_edges")
    op.drop_table("graph_edges")
    op.drop_index("idx_document_chunks_section", table_name="document_chunks")
    op.drop_index("idx_document_chunks_scope", table_name="document_chunks")
    op.drop_index("idx_document_chunks_doc_revision", table_name="document_chunks")
    op.drop_index("idx_document_chunks_chunk_id", table_name="document_chunks")
    op.drop_table("document_chunks")
    op.drop_index(
        "uq_retrieval_hit_stats_chunk_key",
        table_name="retrieval_hit_stats",
        postgresql_where=sa.text("chunk_id IS NOT NULL"),
    )
    op.drop_index(
        "uq_retrieval_hit_stats_document_key",
        table_name="retrieval_hit_stats",
        postgresql_where=sa.text("chunk_id IS NULL"),
    )
    op.drop_index(
        "idx_retrieval_hit_stats_scope_kind", table_name="retrieval_hit_stats"
    )
    op.drop_index("idx_retrieval_hit_stats_document", table_name="retrieval_hit_stats")
    op.drop_index("idx_retrieval_hit_stats_chunk", table_name="retrieval_hit_stats")
    op.drop_table("retrieval_hit_stats")
    op.drop_index("idx_graph_nodes_scope", table_name="graph_nodes")
    op.drop_index("idx_graph_nodes_ref_section", table_name="graph_nodes")
    op.drop_index("idx_graph_nodes_ref_document", table_name="graph_nodes")
    op.drop_index("idx_graph_nodes_owner_revision", table_name="graph_nodes")
    op.drop_table("graph_nodes")
    op.drop_index("idx_document_sections_scope", table_name="document_sections")
    op.drop_index("idx_document_sections_doc_revision", table_name="document_sections")
    op.drop_table("document_sections")
    op.drop_index("idx_documents_user_namespace_status", table_name="documents")
    op.drop_index("idx_documents_current_job_result", table_name="documents")
    op.drop_table("documents")
    op.execute("DROP INDEX IF EXISTS uq_jobs_user_active_document")
    # ### end Alembic commands ###
