import asyncio
from typing import Iterable

from sqlalchemy import inspect, text
from sqlalchemy.ext.asyncio import create_async_engine

from app.core.config import settings
from app.db.base import Base


CRITICAL_TABLE_COLUMNS: dict[str, set[str]] = {
    "studies": {
        "is_locked",
        "visit_schedule",
        "enrollment_monthly_goal_note",
        "enrollment_stage_breakdown",
    },
    "subjects": {
        "baseline_date",
        "actual_medication_count",
    },
    "monitoring_visit_issues": {
        "site_id",
        "severity",
        "mark",
        "visit_cycle",
        "center_query",
        "center_latest_reply",
        "rectification_completed",
    },
}


async def check() -> int:
    engine = create_async_engine(settings.DATABASE_URL)
    try:
        async with engine.connect() as conn:
            def inspect_schema(sync_conn):
                inspector = inspect(sync_conn)
                errors: list[str] = []

                if not inspector.has_table("alembic_version"):
                    errors.append("missing alembic_version table")
                else:
                    version_rows = sync_conn.execute(text("select version_num from alembic_version")).fetchall()
                    if not version_rows:
                        errors.append("alembic_version table is empty")

                for table_name, expected_columns in CRITICAL_TABLE_COLUMNS.items():
                    if not inspector.has_table(table_name):
                        errors.append(f"missing table: {table_name}")
                        continue
                    existing = {col["name"] for col in inspector.get_columns(table_name)}
                    missing = sorted(expected_columns - existing)
                    if missing:
                        errors.append(f"{table_name} missing columns: {', '.join(missing)}")

                return errors

            errors = await conn.run_sync(inspect_schema)
        if errors:
            for item in errors:
                print(f"ERROR: {item}")
            return 1
        print("migration state OK")
        return 0
    finally:
        await engine.dispose()


def main() -> None:
    raise SystemExit(asyncio.run(check()))


if __name__ == "__main__":
    main()