from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import declarative_base, sessionmaker

DATABASE_URL = "sqlite:///./agent_arena.db"

engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def create_all():
    from models import Task, Agent, Run  # noqa: F401
    Base.metadata.create_all(bind=engine)
    _ensure_run_schema()


def _ensure_run_schema():
    expected = {
        "functional_success": "INTEGER DEFAULT 0",
        "test_pass_rate": "FLOAT",
        "alignment_score_llm": "FLOAT",
        "alignment_score_human": "FLOAT",
        "alignment_rubric_version": "VARCHAR(100)",
        "quality_score": "FLOAT",
        "agent_duration_s": "FLOAT DEFAULT 0.0",
        "end_to_end_duration_s": "FLOAT DEFAULT 0.0",
        "timeout_flag": "INTEGER DEFAULT 0",
        "tool_failure_flag": "INTEGER DEFAULT 0",
        "invalid_patch_flag": "INTEGER DEFAULT 0",
        "task_family": "VARCHAR(100)",
        "seed": "VARCHAR(100)",
        "replicate_index": "INTEGER DEFAULT 0",
        "experiment_id": "VARCHAR(100)",
        "arm_id": "VARCHAR(100)",
    }

    inspector = inspect(engine)
    columns = {col["name"] for col in inspector.get_columns("runs")}
    missing = [(name, ddl) for name, ddl in expected.items() if name not in columns]
    if not missing:
        return

    with engine.begin() as conn:
        for name, ddl in missing:
            conn.execute(text(f"ALTER TABLE runs ADD COLUMN {name} {ddl}"))
        conn.execute(text(
            "UPDATE runs SET functional_success = CASE "
            "WHEN functional_success IS NULL OR functional_success = 0 THEN "
            "CASE WHEN status = 'success' THEN 1 ELSE 0 END "
            "ELSE functional_success END"
        ))
        conn.execute(text(
            "UPDATE runs SET agent_duration_s = COALESCE(agent_duration_s, duration_s), "
            "end_to_end_duration_s = COALESCE(end_to_end_duration_s, duration_s), "
            "task_family = COALESCE(task_family, ''), "
            "replicate_index = COALESCE(replicate_index, 0)"
        ))