# ruff: noqa
"""add generic enrichment type

Revision ID: 19f8c7faf8b9
Revises: f9e5ef5e688f
Create Date: 2025-09-29 21:38:19.093821

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa

from kodit.domain.enrichments.development.development import ENRICHMENT_TYPE_DEVELOPMENT
from kodit.domain.enrichments.development.snippet.snippet import (
    ENRICHMENT_SUBTYPE_SNIPPET_SUMMARY,
)


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


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "enrichments_v2",
        sa.Column("content", sa.UnicodeText(), nullable=False),
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("type", sa.String(length=255), nullable=False),
        sa.Column("subtype", sa.String(length=255), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        "idx_type_subtype", "enrichments_v2", ["type", "subtype"], unique=False
    )
    op.create_index(
        op.f("ix_enrichments_v2_type"), "enrichments_v2", ["type"], unique=False
    )
    op.create_index(
        op.f("ix_enrichments_v2_subtype"), "enrichments_v2", ["subtype"], unique=False
    )
    op.create_table(
        "enrichment_associations",
        sa.Column("enrichment_id", sa.Integer(), nullable=False),
        sa.Column("entity_type", sa.String(length=50), nullable=False),
        sa.Column("entity_id", sa.String(length=255), nullable=False),
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["enrichment_id"], ["enrichments_v2.id"], ondelete="CASCADE"
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "entity_type", "entity_id", "enrichment_id", name="uix_entity_enrichment"
        ),
        sqlite_autoincrement=True,
    )
    op.create_index(
        "idx_entity_lookup",
        "enrichment_associations",
        ["entity_type", "entity_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_enrichment_associations_enrichment_id"),
        "enrichment_associations",
        ["enrichment_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_enrichment_associations_entity_id"),
        "enrichment_associations",
        ["entity_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_enrichment_associations_entity_type"),
        "enrichment_associations",
        ["entity_type"],
        unique=False,
    )

    # Migrate existing enrichments from old table to new structure
    # All old enrichments are snippet enrichments
    connection = op.get_bind()

    # Check if old enrichments table exists
    inspector = sa.inspect(connection)
    if "enrichments" in inspector.get_table_names():
        # Copy enrichments to new table
        old_enrichments = connection.execute(
            sa.text(
                "SELECT id, content, created_at, updated_at, snippet_sha FROM enrichments"
            )
        ).fetchall()

        for old_enrichment in old_enrichments:
            # Check if database supports RETURNING (PostgreSQL) or use lastrowid (SQLite)
            dialect_name = connection.dialect.name

            if dialect_name == "postgresql":
                result = connection.execute(
                    sa.text(
                        "INSERT INTO enrichments_v2 (type, subtype, content, created_at, updated_at) "
                        "VALUES (:type, :subtype, :content, :created_at, :updated_at) "
                        "RETURNING id"
                    ),
                    {
                        "type": ENRICHMENT_TYPE_DEVELOPMENT,
                        "subtype": ENRICHMENT_SUBTYPE_SNIPPET_SUMMARY,
                        "content": old_enrichment[1],
                        "created_at": old_enrichment[2],
                        "updated_at": old_enrichment[3],
                    },
                )
                row = result.fetchone()
                if row is None:
                    raise RuntimeError("Failed to insert enrichment")
                new_enrichment_id = row[0]
            else:
                # SQLite and other databases
                result = connection.execute(
                    sa.text(
                        "INSERT INTO enrichments_v2 (type, subtype, content, created_at, updated_at) "
                        "VALUES (:type, :subtype, :content, :created_at, :updated_at)"
                    ),
                    {
                        "type": ENRICHMENT_TYPE_DEVELOPMENT,
                        "subtype": ENRICHMENT_SUBTYPE_SNIPPET_SUMMARY,
                        "content": old_enrichment[1],
                        "created_at": old_enrichment[2],
                        "updated_at": old_enrichment[3],
                    },
                )
                new_enrichment_id = result.lastrowid

            # Insert association (snippet_v2 is the entity type for snippets)
            connection.execute(
                sa.text(
                    "INSERT INTO enrichment_associations "
                    "(enrichment_id, entity_type, entity_id, created_at, updated_at) "
                    "VALUES (:enrichment_id, :entity_type, :entity_id, :created_at, :updated_at)"
                ),
                {
                    "enrichment_id": new_enrichment_id,
                    "entity_type": "snippet_v2",
                    "entity_id": old_enrichment[4],  # snippet_sha
                    "created_at": old_enrichment[2],
                    "updated_at": old_enrichment[3],
                },
            )

        # Drop old enrichments table
        op.drop_table("enrichments")

    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    # Recreate old enrichments table
    op.create_table(
        "enrichments",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
        ),
        sa.Column("snippet_sha", sa.String(length=64), nullable=False),
        sa.Column("type", sa.String(length=50), nullable=False),
        sa.Column("content", sa.UnicodeText(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
        sa.ForeignKeyConstraint(["snippet_sha"], ["snippets_v2.sha"]),
    )
    op.create_index(
        op.f("ix_enrichments_snippet_sha"), "enrichments", ["snippet_sha"], unique=False
    )
    op.create_index(op.f("ix_enrichments_type"), "enrichments", ["type"], unique=False)

    # Migrate data back from new structure to old table
    connection = op.get_bind()

    # Get enrichments with their associations
    enrichments_with_associations = connection.execute(
        sa.text(
            "SELECT e.id, e.content, e.created_at, e.updated_at, "
            "a.entity_id, a.entity_type "
            "FROM enrichments_v2 e "
            "JOIN enrichment_associations a ON e.id = a.enrichment_id "
            "WHERE a.entity_type = 'snippet_v2'"
        )
    ).fetchall()

    for enrichment in enrichments_with_associations:
        connection.execute(
            sa.text(
                "INSERT INTO enrichments "
                "(content, created_at, updated_at, snippet_sha, type) "
                "VALUES (:content, :created_at, :updated_at, :snippet_sha, :type)"
            ),
            {
                "content": enrichment[1],
                "created_at": enrichment[2],
                "updated_at": enrichment[3],
                "snippet_sha": enrichment[4],
                "type": "summarization",  # Default to summarization type
            },
        )

    op.drop_index(
        op.f("ix_enrichment_associations_entity_type"),
        table_name="enrichment_associations",
    )
    op.drop_index(
        op.f("ix_enrichment_associations_entity_id"),
        table_name="enrichment_associations",
    )
    op.drop_index(
        op.f("ix_enrichment_associations_enrichment_id"),
        table_name="enrichment_associations",
    )
    op.drop_index("idx_entity_lookup", table_name="enrichment_associations")
    op.drop_table("enrichment_associations")
    op.drop_index(op.f("ix_enrichments_v2_subtype"), table_name="enrichments_v2")
    op.drop_index(op.f("ix_enrichments_v2_type"), table_name="enrichments_v2")
    op.drop_index("idx_type_subtype", table_name="enrichments_v2")
    op.drop_table("enrichments_v2")
    # ### end Alembic commands ###
