from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path
from typing import Any


SCHEMA_VERSION = 18
SCHEMA_DIR = Path(__file__).with_name("schema")
SCHEMA_PATH = SCHEMA_DIR / "tidb.sql"
MIGRATIONS_DIR = SCHEMA_DIR / "migrations"


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def ensure_tidb_schema(conn: Any) -> int:
    migration_table = conn.execute(
        """
        SELECT COUNT(*) AS table_count
        FROM information_schema.tables
        WHERE table_schema = DATABASE() AND table_name = 'schema_migrations'
        """
    ).fetchone()
    if not migration_table or int(migration_table["table_count"]) == 0:
        conn.executescript(SCHEMA_PATH.read_text(encoding="utf-8"))
        conn.execute(
            "INSERT INTO schema_migrations (component, version, updated_at) "
            "VALUES (?, ?, ?)",
            ("sql_catalog", SCHEMA_VERSION, now_iso()),
        )
        return SCHEMA_VERSION

    row = conn.execute(
        "SELECT version FROM schema_migrations WHERE component = ?",
        ("sql_catalog",),
    ).fetchone()
    if not row:
        raise RuntimeError(
            "TiDB schema_migrations is missing the sql_catalog component"
        )
    current_version = int(row["version"])
    if current_version > SCHEMA_VERSION:
        raise RuntimeError(
            "TiDB schema is newer than this application: "
            f"database={current_version}, application={SCHEMA_VERSION}"
        )

    for target_version in range(current_version + 1, SCHEMA_VERSION + 1):
        matches = sorted(MIGRATIONS_DIR.glob(f"{target_version:04d}_*.sql"))
        if len(matches) != 1:
            raise RuntimeError(
                f"Expected exactly one TiDB migration for version {target_version}"
            )
        conn.executescript(matches[0].read_text(encoding="utf-8"))
        conn.execute(
            "UPDATE schema_migrations SET version = ?, updated_at = ? "
            "WHERE component = ?",
            (target_version, now_iso(), "sql_catalog"),
        )
        current_version = target_version
    return current_version