from pathlib import Path
from tempfile import TemporaryDirectory

from alembic.config import Config
from sqlalchemy import create_engine, inspect

from alembic import command
from modelforge_api.persistence.models import Base


def test_operator_benchmark_migration_upgrades_and_downgrades_sqlite() -> None:
    with TemporaryDirectory() as temp:
        database_path = Path(temp) / "benchmark-migration.db"
        url = f"sqlite:///{database_path.as_posix()}"
        engine = create_engine(url)
        Base.metadata.create_all(engine)
        with engine.begin() as connection:
            connection.exec_driver_sql("DROP TABLE benchmark_runs")
            connection.exec_driver_sql(
                """
                CREATE TABLE benchmark_runs (
                    id CHAR(32) NOT NULL,
                    deployment_id CHAR(32) NOT NULL,
                    suite_id CHAR(32) NOT NULL,
                    model_revision_sha VARCHAR(64) NOT NULL,
                    artifact_sha256 VARCHAR(64) NOT NULL,
                    runtime_type VARCHAR(32) NOT NULL,
                    runtime_version VARCHAR(128) NOT NULL,
                    runtime_image_digest VARCHAR(71),
                    launch_arguments JSON NOT NULL,
                    suite_revision VARCHAR(128) NOT NULL,
                    dataset_revision VARCHAR(255) NOT NULL,
                    environment_fingerprint JSON NOT NULL,
                    environment_digest VARCHAR(64) NOT NULL,
                    seed INTEGER,
                    results JSON NOT NULL,
                    status VARCHAR(32) NOT NULL,
                    comparable BOOLEAN NOT NULL,
                    stale_reason TEXT,
                    started_at DATETIME NOT NULL,
                    completed_at DATETIME,
                    PRIMARY KEY (id),
                    FOREIGN KEY(deployment_id) REFERENCES deployments (id) ON DELETE RESTRICT,
                    FOREIGN KEY(suite_id) REFERENCES benchmark_suites (id) ON DELETE RESTRICT
                )
                """
            )
            connection.exec_driver_sql(
                "CREATE INDEX ix_benchmark_runs_deployment_id ON benchmark_runs (deployment_id)"
            )
            connection.exec_driver_sql(
                "CREATE INDEX ix_benchmark_runs_suite_id ON benchmark_runs (suite_id)"
            )
            connection.exec_driver_sql(
                "CREATE INDEX ix_benchmark_runs_artifact_sha256 ON benchmark_runs (artifact_sha256)"
            )
            connection.exec_driver_sql(
                "CREATE INDEX ix_benchmark_runs_environment_digest ON benchmark_runs (environment_digest)"
            )
            connection.exec_driver_sql(
                "CREATE INDEX ix_benchmark_runs_deployment_completed "
                "ON benchmark_runs (deployment_id, completed_at)"
            )

        config = Config(str(Path(__file__).parents[1] / "alembic.ini"))
        config.set_main_option("sqlalchemy.url", url)
        command.stamp(config, "20260924_0038")
        command.upgrade(config, "head")

        upgraded = inspect(engine)
        columns = {column["name"]: column for column in upgraded.get_columns("benchmark_runs")}
        assert columns["deployment_id"]["nullable"] is True
        assert "capability_deployment_id" in columns
        indexes = {index["name"] for index in upgraded.get_indexes("benchmark_runs")}
        assert "ix_benchmark_runs_actor_subject" in indexes
        assert "ix_benchmark_runs_capability_deployment_completed" in indexes

        command.downgrade(config, "20260924_0038")
        downgraded = {
            column["name"]: column for column in inspect(engine).get_columns("benchmark_runs")
        }
        assert downgraded["deployment_id"]["nullable"] is False
        assert "capability_deployment_id" not in downgraded

        command.upgrade(config, "head")
        assert "capability_deployment_id" in {
            column["name"] for column in inspect(engine).get_columns("benchmark_runs")
        }
        engine.dispose()