"""Apply db/migrations/002_api_keys_field_kind.sql to the Normbase Postgres database once.

Usage (from repo root):
  python normbase/scripts/apply_migration_002_api_keys_field.py

Uses NORMBASE_DATABASE_URL or DATABASE_URL from the environment (.env loaded if present).
"""
from __future__ import annotations

import asyncio
import os
import re
import sys
from pathlib import Path

try:
    import asyncpg
except ImportError as e:
    raise SystemExit("asyncpg is required. Install project dependencies first.") from e


def _dsn() -> str:
    try:
        from dotenv import load_dotenv

        load_dotenv()
    except ImportError:
        pass
    raw = (os.getenv("NORMBASE_DATABASE_URL") or os.getenv("DATABASE_URL") or "").strip()
    if not raw:
        raise SystemExit("Set NORMBASE_DATABASE_URL or DATABASE_URL.")
    # SQLAlchemy / async URL -> asyncpg
    return re.sub(r"^postgresql\+asyncpg", "postgresql", raw, flags=re.I)


# Keep in sync with normbase/db/migrations/002_api_keys_field_kind.sql
_STATEMENTS = [
    "ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS key_kind VARCHAR(20) NOT NULL DEFAULT 'agent'",
    "ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS field_id UUID REFERENCES fields(id) ON DELETE CASCADE",
    "UPDATE api_keys SET key_kind = 'agent' WHERE key_kind IS NULL OR key_kind = ''",
    "ALTER TABLE api_keys DROP CONSTRAINT IF EXISTS ck_api_keys_kind_target",
    """ALTER TABLE api_keys ADD CONSTRAINT ck_api_keys_kind_target CHECK (
    (key_kind = 'agent' AND agent_id IS NOT NULL AND field_id IS NULL)
    OR (key_kind = 'field' AND field_id IS NOT NULL AND agent_id IS NULL)
)""",
    "ALTER TABLE api_keys DROP CONSTRAINT IF EXISTS ck_api_keys_key_kind_value",
    "ALTER TABLE api_keys ADD CONSTRAINT ck_api_keys_key_kind_value CHECK (key_kind IN ('agent', 'field'))",
]


async def _run() -> None:
    conn = await asyncpg.connect(_dsn())
    try:
        for i, sql in enumerate(_STATEMENTS, 1):
            await conn.execute(sql)
            print(f"  [{i}/{len(_STATEMENTS)}] ok")
        print("Migration 002_api_keys_field_kind applied successfully.")
    finally:
        await conn.close()


def main() -> None:
    repo_root = Path(__file__).resolve().parents[2]
    os.chdir(repo_root)
    asyncio.run(_run())


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(f"Migration failed: {exc}", file=sys.stderr)
        sys.exit(1)
