"""create_core_task_runs and backfill code_pedimento_regimens

Revision ID: c1a2b3d4e5f6
Revises: bccb7f8986c7
Create Date: 2026-03-24 11:30:00.000000
"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from api.v1.modules.public.reference_data.code_pedimento_regimens.seed import (
    seed as code_pedimento_regimens_seed_full,
)
from api.v1.modules.public.reference_data.pedimento_codes.seed import (
    seed as pedimento_codes_seed,
)


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


def _format_value(val: str | None) -> str:
    if val is None or str(val).strip() == "" or str(val).upper() == "NONE":
        return "NULL"
    return f"'{str(val).replace(chr(39), chr(39) * 2)}'"


def upgrade() -> None:
    op.create_table(
        "task_runs",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("task_id", sa.String(length=255), nullable=False),
        sa.Column("tenant_id", sa.Integer(), nullable=False),
        sa.Column("company_id", sa.Integer(), nullable=True),
        sa.Column("requested_by_user", sa.String(length=255), nullable=True),
        sa.Column("task_name", sa.String(length=255), nullable=False),
        sa.Column("task_group", sa.String(length=100), nullable=False),
        sa.Column("task_origin", sa.String(length=255), nullable=True),
        sa.Column("status", sa.String(length=20), nullable=False),
        sa.Column("celery_state_raw", sa.String(length=30), nullable=False),
        sa.Column("progress_current", sa.Integer(), nullable=True),
        sa.Column("progress_total", sa.Integer(), nullable=True),
        sa.Column("progress_percent", sa.Float(), nullable=True),
        sa.Column("progress_message", sa.String(length=500), nullable=True),
        sa.Column("retries", sa.Integer(), nullable=True),
        sa.Column("exception_type", sa.String(length=255), nullable=True),
        sa.Column("exception_message", sa.Text(), nullable=True),
        sa.Column("traceback_excerpt", sa.Text(), nullable=True),
        sa.Column("result_summary", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.Column("meta_payload", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(["company_id"], ["a76.company.id"]),
        sa.ForeignKeyConstraint(["tenant_id"], ["core.tenants.id"]),
        sa.PrimaryKeyConstraint("id"),
        schema="core",
    )
    op.create_index("ix_core_task_runs_task_id", "task_runs", ["task_id"], unique=True, schema="core")
    op.create_index("ix_core_task_runs_tenant_updated", "task_runs", ["tenant_id", "updated_at"], unique=False, schema="core")
    op.create_index(
        "ix_core_task_runs_tenant_status_updated",
        "task_runs",
        ["tenant_id", "status", "updated_at"],
        unique=False,
        schema="core",
    )
    op.create_index(
        "ix_core_task_runs_tenant_group_updated",
        "task_runs",
        ["tenant_id", "task_group", "updated_at"],
        unique=False,
        schema="core",
    )
    op.create_index(
        "ix_core_task_runs_tenant_company_updated",
        "task_runs",
        ["tenant_id", "company_id", "updated_at"],
        unique=False,
        schema="core",
    )

    pedimento_desc_by_code = {code: desc for code, desc in pedimento_codes_seed}
    required_pedimento_codes = sorted(
        {pedimento_code for pedimento_code, _regimen_code, _type_code in code_pedimento_regimens_seed_full}
    )

    values_missing_codes = ", ".join(
        [
            f"({_format_value(code)}, {_format_value(pedimento_desc_by_code.get(code, f'AUTO-GENERATED FOR FK ({code})'))})"
            for code in required_pedimento_codes
        ]
    )
    op.execute(
        f"""
        INSERT INTO public.pedimento_codes (code, description)
        VALUES {values_missing_codes}
        ON CONFLICT (code) DO NOTHING;
    """
    )

    values_relations = ", ".join(
        [
            f"({_format_value(pedimento_code)}, {_format_value(regimen_code)}, {_format_value(type_code)})"
            for pedimento_code, regimen_code, type_code in code_pedimento_regimens_seed_full
        ]
    )
    op.execute(
        f"""
        INSERT INTO public.code_pedimento_regimens (pedimento_code, regimen_code, type_code)
        SELECT src.pedimento_code, src.regimen_code, src.type_code
        FROM (VALUES {values_relations}) AS src(pedimento_code, regimen_code, type_code)
        JOIN public.pedimento_codes pc
            ON pc.code = src.pedimento_code
        JOIN public.pedimento_regimens pr
            ON pr.code = src.regimen_code
        WHERE NOT EXISTS (
            SELECT 1
            FROM public.code_pedimento_regimens cpr
            WHERE cpr.pedimento_code = src.pedimento_code
              AND cpr.regimen_code = src.regimen_code
              AND COALESCE(cpr.type_code, '') = COALESCE(src.type_code, '')
        );
    """
    )


def downgrade() -> None:
    op.drop_index("ix_core_task_runs_tenant_company_updated", table_name="task_runs", schema="core")
    op.drop_index("ix_core_task_runs_tenant_group_updated", table_name="task_runs", schema="core")
    op.drop_index("ix_core_task_runs_tenant_status_updated", table_name="task_runs", schema="core")
    op.drop_index("ix_core_task_runs_tenant_updated", table_name="task_runs", schema="core")
    op.drop_index("ix_core_task_runs_task_id", table_name="task_runs", schema="core")
    op.drop_table("task_runs", schema="core")
