"""
Simple migration runner using asyncpg.
Tracks applied migrations in a _migrations table.

Usage:
    DATABASE_URL=postgresql://user:pass@localhost/db uv run python migrations/migrate.py apply
    DATABASE_URL=postgresql://user:pass@localhost/db uv run python migrations/migrate.py status
"""

import asyncio
import os
import sys
from pathlib import Path

import asyncpg

MIGRATIONS_DIR = Path(__file__).parent


async def ensure_migrations_table(conn: asyncpg.Connection) -> None:
    """Create the migrations tracking table if it doesn't exist."""
    await conn.execute("""
        CREATE TABLE IF NOT EXISTS _migrations (
            id SERIAL PRIMARY KEY,
            name TEXT NOT NULL UNIQUE,
            applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
        )
    """)


async def get_applied_migrations(conn: asyncpg.Connection) -> set[str]:
    """Get the set of already applied migration names."""
    rows = await conn.fetch("SELECT name FROM _migrations")
    return {row["name"] for row in rows}


async def get_pending_migrations(conn: asyncpg.Connection) -> list[Path]:
    """Get list of migration files that haven't been applied yet."""
    applied = await get_applied_migrations(conn)
    sql_files = sorted(MIGRATIONS_DIR.glob("*.sql"))
    return [f for f in sql_files if f.name not in applied]


async def apply_migration(conn: asyncpg.Connection, migration_file: Path) -> None:
    """Apply a single migration file within a transaction."""
    sql = migration_file.read_text()
    async with conn.transaction():
        await conn.execute(sql)
        await conn.execute(
            "INSERT INTO _migrations (name) VALUES ($1)",
            migration_file.name
        )
    print(f"Applied: {migration_file.name}")


async def migrate(database_url: str) -> None:
    """Apply all pending migrations."""
    conn = await asyncpg.connect(database_url)
    try:
        await ensure_migrations_table(conn)
        pending = await get_pending_migrations(conn)

        if not pending:
            print("No pending migrations.")
            return

        for migration_file in pending:
            await apply_migration(conn, migration_file)

        print(f"Applied {len(pending)} migration(s).")
    finally:
        await conn.close()


async def status(database_url: str) -> None:
    """Show migration status."""
    conn = await asyncpg.connect(database_url)
    try:
        await ensure_migrations_table(conn)
        applied = await get_applied_migrations(conn)
        pending = await get_pending_migrations(conn)

        print("Applied migrations:")
        for name in sorted(applied):
            print(f"  [x] {name}")

        print("\nPending migrations:")
        for f in pending:
            print(f"  [ ] {f.name}")

        if not applied and not pending:
            print("  (none)")
    finally:
        await conn.close()


def main() -> None:
    database_url = os.environ.get("DATABASE_URL")
    if not database_url:
        print("Error: DATABASE_URL environment variable is required")
        sys.exit(1)

    if len(sys.argv) < 2:
        print("Usage: python migrate.py [apply|status]")
        sys.exit(1)

    command = sys.argv[1]
    if command == "apply":
        asyncio.run(migrate(database_url))
    elif command == "status":
        asyncio.run(status(database_url))
    else:
        print(f"Unknown command: {command}")
        print("Usage: python migrate.py [apply|status]")
        sys.exit(1)


if __name__ == "__main__":
    main()