"""jsonb schema validation case fields

Revision ID: 3f1bacf5a8a9
Revises: b873c83b8e0a
Create Date: 2025-11-25 20:33:14.153722

"""

import json
from collections.abc import Sequence

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op
from tracecat.identifiers.workflow import WorkspaceUUID

# revision identifiers, used by Alembic.
revision: str = "3f1bacf5a8a9"
down_revision: str | None = "b873c83b8e0a"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

# Reserved columns - not custom fields
RESERVED_COLUMNS = {"id", "case_id", "created_at", "updated_at", "owner_id"}

# Map SQLAlchemy types to logical types
SA_TYPE_MAP = {
    sa.TEXT: "TEXT",
    sa.String: "TEXT",
    sa.VARCHAR: "TEXT",
    sa.Integer: "INTEGER",
    sa.BigInteger: "INTEGER",
    sa.SmallInteger: "INTEGER",
    sa.Numeric: "NUMERIC",
    sa.Float: "NUMERIC",
    sa.Boolean: "BOOLEAN",
    sa.TIMESTAMP: "TIMESTAMPTZ",
    sa.DateTime: "TIMESTAMPTZ",
    postgresql.JSONB: "JSONB",
    postgresql.JSON: "JSONB",
}


def _get_logical_type(sa_type: sa.types.TypeEngine) -> str:
    """Map SQLAlchemy type to our logical type."""
    for sa_cls, logical in SA_TYPE_MAP.items():
        if isinstance(sa_type, sa_cls):
            return logical
    return "TEXT"


def _populate_schema_from_workspace_tables(connection: sa.Connection) -> None:
    """Inspect dynamic workspace tables and populate the schema column."""
    inspector = sa.inspect(connection)

    # Get all custom_fields_ schemas
    schemas = [
        s for s in inspector.get_schema_names() if s.startswith("custom_fields_")
    ]

    for schema_name in schemas:
        # Extract workspace short ID from schema name (custom_fields_XXXXX)
        workspace_short_id = schema_name.replace("custom_fields_", "")

        # Check if case_fields table exists in this schema
        if not inspector.has_table("case_fields", schema=schema_name):
            continue

        # Reflect columns using SQLAlchemy inspector
        columns = inspector.get_columns("case_fields", schema=schema_name)

        # Build schema dict for custom columns only
        field_schema = {}
        for col in columns:
            col_name = col["name"]
            if col_name in RESERVED_COLUMNS:
                continue
            field_schema[col_name] = {"type": _get_logical_type(col["type"])}

        if not field_schema:
            continue

        # Find matching case_fields row by workspace short ID
        workspace_result = connection.execute(
            sa.text("SELECT id, owner_id FROM case_fields")
        )

        for cf_id, owner_id in workspace_result:
            short_id = WorkspaceUUID.new(owner_id).short()
            if short_id == workspace_short_id:
                connection.execute(
                    sa.text("""
                        UPDATE case_fields
                        SET schema = :schema
                        WHERE id = :cf_id
                    """),
                    {"schema": json.dumps(field_schema), "cf_id": cf_id},
                )
                break


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###

    # Add schema column with default empty object
    op.add_column(
        "case_fields",
        sa.Column(
            "schema",
            postgresql.JSONB(astext_type=sa.Text()),
            nullable=False,
            server_default="{}",
        ),
    )

    # Deduplicate case_fields - keep only one row per owner_id (oldest)
    op.execute("""
        DELETE FROM case_fields
        WHERE id NOT IN (
            SELECT DISTINCT ON (owner_id) id
            FROM case_fields
            ORDER BY owner_id, created_at ASC
        )
    """)

    # Populate schema from existing dynamic workspace tables
    connection = op.get_bind()
    _populate_schema_from_workspace_tables(connection)

    op.drop_constraint("case_fields_case_id_key", "case_fields", type_="unique")
    op.create_unique_constraint("uq_case_fields_owner", "case_fields", ["owner_id"])
    op.drop_constraint("case_fields_case_id_fkey", "case_fields", type_="foreignkey")
    op.drop_column("case_fields", "case_id")

    # Remove server default after migration
    op.alter_column("case_fields", "schema", server_default=None)
    # ### end Alembic commands ###


def _repopulate_case_fields_from_workspace_tables(connection: sa.Connection) -> None:
    """Re-create case_fields rows from dynamic workspace tables for downgrade."""
    inspector = sa.inspect(connection)

    # Get all custom_fields_ schemas
    schemas = [
        s for s in inspector.get_schema_names() if s.startswith("custom_fields_")
    ]

    for schema_name in schemas:
        workspace_short_id = schema_name.replace("custom_fields_", "")

        if not inspector.has_table("case_fields", schema=schema_name):
            continue

        # Find the workspace owner_id by matching short ID
        workspace_result = connection.execute(
            sa.text("SELECT id, owner_id FROM case_fields")
        )

        owner_id = None
        existing_cf_id = None
        for cf_id, oid in workspace_result:
            short_id = WorkspaceUUID.new(oid).short()
            if short_id == workspace_short_id:
                owner_id = oid
                existing_cf_id = cf_id
                break

        if owner_id is None:
            continue

        # Get all rows from the dynamic workspace table
        rows = connection.execute(
            sa.text(f'SELECT id, case_id FROM "{schema_name}".case_fields')
        ).fetchall()

        # For each row in the dynamic table, create/update a case_fields row
        for i, (_, case_id) in enumerate(rows):
            if i == 0 and existing_cf_id:
                # Update the existing case_fields row with the first case_id
                connection.execute(
                    sa.text("""
                        UPDATE case_fields
                        SET case_id = :case_id
                        WHERE id = :cf_id
                    """),
                    {"case_id": case_id, "cf_id": existing_cf_id},
                )
            else:
                # Insert new case_fields rows for additional cases
                connection.execute(
                    sa.text("""
                        INSERT INTO case_fields (id, owner_id, case_id, created_at, updated_at)
                        VALUES (gen_random_uuid(), :owner_id, :case_id, now(), now())
                    """),
                    {"owner_id": owner_id, "case_id": case_id},
                )


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###

    # Drop schema column first (before inserting new rows)
    op.drop_column("case_fields", "schema")

    # Drop the unique constraint on owner_id (allows multiple rows per workspace)
    op.drop_constraint("uq_case_fields_owner", "case_fields", type_="unique")

    # Add case_id column as nullable
    op.add_column(
        "case_fields",
        sa.Column("case_id", sa.UUID(), autoincrement=False, nullable=True),
    )

    # Repopulate case_fields from dynamic workspace tables
    connection = op.get_bind()
    _repopulate_case_fields_from_workspace_tables(connection)

    # Delete any case_fields rows that couldn't be linked to a case
    op.execute("DELETE FROM case_fields WHERE case_id IS NULL")

    # Now make case_id NOT NULL
    op.alter_column("case_fields", "case_id", nullable=False)

    # Re-create constraints
    op.create_foreign_key(
        "case_fields_case_id_fkey",
        "case_fields",
        "cases",
        ["case_id"],
        ["id"],
        ondelete="CASCADE",
    )
    op.create_unique_constraint(
        "case_fields_case_id_key",
        "case_fields",
        ["case_id"],
        postgresql_nulls_not_distinct=False,
    )
    # ### end Alembic commands ###
