from __future__ import annotations

import importlib.util
from pathlib import Path

import sqlalchemy as sa
from alembic.migration import MigrationContext
from alembic.operations import Operations


def _load_migration_module():
    path = (
        Path(__file__).resolve().parents[1]
        / "alembic"
        / "versions"
        / "20260327_add_reply_iteration_training_tables.py"
    )
    spec = importlib.util.spec_from_file_location("reply_iteration_training_tables_migration", path)
    module = importlib.util.module_from_spec(spec)
    assert spec is not None
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def _create_base_tables(metadata: sa.MetaData) -> None:
    sa.Table(
        "feedback_center_records",
        metadata,
        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column("tenant_id", sa.BigInteger(), nullable=False),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.Column("updated_at", sa.TIMESTAMP(), nullable=False),
    )
    sa.Table(
        "qa_records",
        metadata,
        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column("tenant_id", sa.BigInteger(), nullable=False),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.Column("updated_at", sa.TIMESTAMP(), nullable=False),
    )
    sa.Table(
        "experiments",
        metadata,
        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
        sa.Column("tenant_id", sa.BigInteger(), nullable=False),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.Column("updated_at", sa.TIMESTAMP(), nullable=False),
        sa.Column("title", sa.String(length=255), nullable=True),
        sa.Column("hypothesis", sa.Text(), nullable=True),
        sa.Column("status", sa.String(length=50), nullable=False),
        sa.Column("target_metric", sa.String(length=50), nullable=True),
        sa.Column("current_score", sa.Float(), nullable=True),
        sa.Column("linked_context_nodes", sa.JSON(), nullable=True),
        sa.Column("project_id", sa.BigInteger(), nullable=True),
    )


def test_reply_iteration_training_migration_adds_tables_and_experiment_columns() -> None:
    migration = _load_migration_module()
    engine = sa.create_engine("sqlite:///:memory:")
    metadata = sa.MetaData()
    _create_base_tables(metadata)
    metadata.create_all(engine)

    with engine.begin() as conn:
        context = MigrationContext.configure(conn)
        migration.op = Operations(context)
        migration.upgrade()

        inspector = sa.inspect(conn)
        tables = set(inspector.get_table_names())
        experiment_columns = {column["name"] for column in inspector.get_columns("experiments")}

        assert "quality_assets" in tables
        assert "quality_asset_pairs" in tables
        assert "judge_runs" in tables
        assert "optimizer_candidates" in tables
        assert "prompt_releases" in tables
        assert "training_datasets" in tables
        assert "training_dataset_items" in tables
        assert "training_jobs" in tables
        assert "training_job_runs" in tables
        assert "training_evaluations" in tables
        assert "trained_model_artifacts" in tables
        assert "trained_model_deployments" in tables
        assert "experiment_type" in experiment_columns
        assert "baseline_release_id" in experiment_columns
        assert "candidate_release_id" in experiment_columns
        assert "scope_json" in experiment_columns
        assert "result_summary" in experiment_columns
        assert "decision" in experiment_columns

        migration.downgrade()

        inspector = sa.inspect(conn)
        tables = set(inspector.get_table_names())
        experiment_columns = {column["name"] for column in inspector.get_columns("experiments")}

        assert "quality_assets" not in tables
        assert "training_jobs" not in tables
        assert "experiment_type" not in experiment_columns
        assert "baseline_release_id" not in experiment_columns