"""Add v2 LLM validation models (tables, columns, indexes)

Adds:
  - validation_sources table
  - validation_runs table
  - v2 columns to validation_policies
  - v2 columns + FK + indexes to llm_validation_results
  - index on validation_policies.provider_id
  - optional data backfill: creates ValidationRun placeholders for existing
    ValidationRecord rows that have a policy_id but no run_id

@ADR [MIGRATION-001] Backfill creates one run per distinct policy_id.
@RATIONALE Existing records were created before the ValidationRun model existed.
  Grouping by policy_id avoids creating excessive runs for what is essentially
  one historic execution batch per policy. The run status is set to 'completed'
  to avoid confusing downstream consumers that filter on 'running'.

Revision ID: c9d8e7f6a5b4
Revises: 2df63b7ce038
Create Date: 2026-05-31 12:00:00.000000

"""
from collections.abc import Sequence

from alembic import op
import sqlalchemy as sa
from sqlalchemy import inspect

# revision identifiers, used by Alembic.
revision: str = "c9d8e7f6a5b4"
down_revision: str | Sequence[str] | None = "2df63b7ce038"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def _table_exists(table_name: str) -> bool:
    conn = op.get_bind()
    inspector = inspect(conn)
    return table_name in inspector.get_table_names()


def upgrade() -> None:
    """Upgrade schema to v2 validation models."""
    _create_validation_sources()
    _create_validation_runs()
    _extend_validation_policies()
    _extend_validation_results()
    _backfill_validation_runs()


def downgrade() -> None:
    """Downgrade schema by removing v2 additions in reverse order."""
    # Drop FK constraints from llm_validation_results
    op.drop_constraint(
        "fk_llm_validation_results_source_id",
        "llm_validation_results",
        type_="foreignkey",
    )
    op.drop_constraint(
        "fk_llm_validation_results_run_id",
        "llm_validation_results",
        type_="foreignkey",
    )
    op.drop_index(
        op.f("ix_llm_validation_results_run_id"),
        table_name="llm_validation_results",
    )

    # Remove v2 columns from llm_validation_results
    _v2_result_columns = [
        "screenshot_paths",
        "timings",
        "token_usage",
        "logs_sent_to_llm",
        "tab_screenshots",
        "chart_data_results",
        "dataset_health",
        "execution_path",
        "source_id",
        "run_id",
    ]
    for col in _v2_result_columns:
        op.drop_column("llm_validation_results", col)

    # Remove columns + index from validation_policies
    op.drop_index(
        op.f("ix_validation_policies_provider_id"),
        table_name="validation_policies",
    )
    _v2_policy_columns = [
        "source_snapshot",
        "policy_dashboard_concurrency_limit",
        "llm_batch_size",
        "execute_chart_data",
        "logs_enabled",
        "screenshot_enabled",
        "prompt_template",
    ]
    for col in _v2_policy_columns:
        op.drop_column("validation_policies", col)

    # Drop validation_runs table (index first, then table)
    op.drop_index(
        op.f("ix_validation_runs_policy_id"),
        table_name="validation_runs",
    )
    op.drop_table("validation_runs")

    # Drop validation_sources table (index first, then table)
    op.drop_index(
        op.f("ix_validation_sources_policy_id"),
        table_name="validation_sources",
    )
    op.drop_table("validation_sources")


# ---------------------------------------------------------------------------
# Upgrade helpers (grouped for readability)
# ---------------------------------------------------------------------------


def _create_validation_sources() -> None:
    """Create the validation_sources table (no FK dependencies)."""
    if not _table_exists("validation_policies"):
        return
    op.create_table(
        "validation_sources",
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("policy_id", sa.String(), nullable=False),
        sa.Column("type", sa.String(), nullable=False),
        sa.Column("value", sa.String(), nullable=False),
        sa.Column("parsed_context", sa.JSON(), nullable=True),
        sa.Column(
            "status", sa.String(), nullable=False, server_default="valid"
        ),
        sa.Column("resolved_dashboard_id", sa.String(), nullable=True),
        sa.Column("title", sa.String(), nullable=True),
        sa.Column("last_error", sa.String(), nullable=True),
        sa.Column("last_checked_at", sa.DateTime(), nullable=True),
        sa.Column(
            "created_at",
            sa.DateTime(),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.ForeignKeyConstraint(
            ["policy_id"],
            ["validation_policies.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_validation_sources_policy_id"),
        "validation_sources",
        ["policy_id"],
    )


def _create_validation_runs() -> None:
    """Create the validation_runs table."""
    if not _table_exists("validation_policies"):
        return
    op.create_table(
        "validation_runs",
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("policy_id", sa.String(), nullable=False),
        sa.Column("task_id", sa.String(), nullable=True),
        sa.Column(
            "started_at",
            sa.DateTime(),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.Column("finished_at", sa.DateTime(), nullable=True),
        sa.Column(
            "trigger", sa.String(), nullable=False, server_default="manual"
        ),
        sa.Column(
            "status", sa.String(), nullable=False, server_default="running"
        ),
        sa.Column(
            "dashboard_count",
            sa.Integer(),
            nullable=False,
            server_default="0",
        ),
        sa.Column(
            "pass_count", sa.Integer(), nullable=False, server_default="0"
        ),
        sa.Column(
            "warn_count", sa.Integer(), nullable=False, server_default="0"
        ),
        sa.Column(
            "fail_count", sa.Integer(), nullable=False, server_default="0"
        ),
        sa.Column(
            "unknown_count", sa.Integer(), nullable=False, server_default="0"
        ),
        sa.Column(
            "created_at",
            sa.DateTime(),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.ForeignKeyConstraint(
            ["policy_id"],
            ["validation_policies.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_validation_runs_policy_id"),
        "validation_runs",
        ["policy_id"],
    )


def _extend_validation_policies() -> None:
    """Add v2 columns and index to validation_policies."""
    if not _table_exists("validation_policies"):
        return
    op.add_column(
        "validation_policies",
        sa.Column("prompt_template", sa.Text(), nullable=True),
    )
    op.add_column(
        "validation_policies",
        sa.Column(
            "screenshot_enabled",
            sa.Boolean(),
            nullable=False,
            server_default="true",
        ),
    )
    op.add_column(
        "validation_policies",
        sa.Column(
            "logs_enabled",
            sa.Boolean(),
            nullable=False,
            server_default="true",
        ),
    )
    op.add_column(
        "validation_policies",
        sa.Column(
            "execute_chart_data",
            sa.Boolean(),
            nullable=False,
            server_default="false",
        ),
    )
    op.add_column(
        "validation_policies",
        sa.Column(
            "llm_batch_size",
            sa.Integer(),
            nullable=False,
            server_default="1",
        ),
    )
    op.add_column(
        "validation_policies",
        sa.Column(
            "policy_dashboard_concurrency_limit",
            sa.Integer(),
            nullable=False,
            server_default="3",
        ),
    )
    op.add_column(
        "validation_policies",
        sa.Column("source_snapshot", sa.JSON(), nullable=True),
    )
    op.create_index(
        op.f("ix_validation_policies_provider_id"),
        "validation_policies",
        ["provider_id"],
    )


def _extend_validation_results() -> None:
    """Add v2 columns, FKs, and indexes to llm_validation_results."""
    if not _table_exists("llm_validation_results"):
        return
    op.add_column(
        "llm_validation_results",
        sa.Column(
            "run_id",
            sa.String(),
            sa.ForeignKey("validation_runs.id", name="fk_llm_validation_results_run_id"),
            nullable=True,
        ),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column(
            "source_id",
            sa.String(),
            sa.ForeignKey("validation_sources.id", name="fk_llm_validation_results_source_id"),
            nullable=True,
        ),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column("execution_path", sa.String(), nullable=True),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column("dataset_health", sa.JSON(), nullable=True),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column("chart_data_results", sa.JSON(), nullable=True),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column("tab_screenshots", sa.JSON(), nullable=True),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column("logs_sent_to_llm", sa.JSON(), nullable=True),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column("token_usage", sa.JSON(), nullable=True),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column("timings", sa.JSON(), nullable=True),
    )
    op.add_column(
        "llm_validation_results",
        sa.Column("screenshot_paths", sa.JSON(), nullable=True),
    )
    op.create_index(
        op.f("ix_llm_validation_results_run_id"),
        "llm_validation_results",
        ["run_id"],
    )


def _backfill_validation_runs() -> None:
    """Create placeholder ValidationRun entries for existing records.

    For each distinct policy_id that has ValidationRecord rows without a
    run_id, creates one ValidationRun (status='completed') and links all
    matching records to it.

    Safe no-op when no such rows exist (e.g. fresh database).
    """
    if not _table_exists("llm_validation_results"):
        return
    conn = op.get_bind()

    # Check for existing records that need backfill
    result = conn.execute(
        sa.text(
            "SELECT DISTINCT v.policy_id "
            "FROM llm_validation_results v "
            "WHERE v.policy_id IS NOT NULL "
            "AND v.run_id IS NULL"
        )
    )
    rows = result.fetchall()
    if not rows:
        return

    import uuid

    for (policy_id,) in rows:
        run_id = str(uuid.uuid4())

        conn.execute(
            sa.text(
                "INSERT INTO validation_runs "
                "(id, policy_id, started_at, trigger, status, created_at) "
                "VALUES (:run_id, :policy_id, NOW(), 'manual', 'completed', NOW())"
            ),
            {"run_id": run_id, "policy_id": policy_id},
        )
        conn.execute(
            sa.text(
                "UPDATE llm_validation_results "
                "SET run_id = :run_id "
                "WHERE policy_id = :policy_id AND run_id IS NULL"
            ),
            {"run_id": run_id, "policy_id": policy_id},
        )
