#!/usr/bin/env python3
"""CLI entry point for repo_core migrations.

Applies schema migrations directly against the database using
:class:`~repo_core.migrations.runner.MigrationRunner` — no Alembic required.

Usage::

    # PostgreSQL
    python -m repo_core.migrations upgrade             # apply all pending migrations
    python -m repo_core.migrations downgrade           # roll back the last migration
    python -m repo_core.migrations downgrade <version> # roll back to a specific version
    python -m repo_core.migrations status              # show migration status

    # Doris (warehouse)
    python -m repo_core.migrations doris-upgrade
    python -m repo_core.migrations doris-downgrade [<version>]
    python -m repo_core.migrations doris-status
"""

from __future__ import annotations

import argparse
import asyncio
import os
import sys
from pathlib import Path



def _load_env() -> None:
    """Load .env from the repo_core project root."""
    try:
        from dotenv import load_dotenv
    except ImportError:
        return
    env_path = Path(__file__).resolve().parent.parent / ".env"
    if env_path.exists():
        load_dotenv(env_path)
    if not os.getenv("DATABASE_URL"):
        _build_pg_url_from_parts()


def _build_pg_url_from_parts() -> None:
    user = os.getenv("PG_USER", os.getenv("POSTGRES_USER", ""))
    password = os.getenv("PG_PASSWORD", os.getenv("POSTGRES_PASSWORD", ""))
    host = os.getenv("PG_HOST", os.getenv("POSTGRES_HOST", "localhost"))
    port = os.getenv("PG_PORT", os.getenv("POSTGRES_PORT", "5432"))
    database = os.getenv("PG_DATABASE", os.getenv("POSTGRES_DB", ""))
    if user and database:
        auth = f"{user}:{password}@" if password else f"{user}@"
        os.environ["DATABASE_URL"] = f"postgresql+asyncpg://{auth}{host}:{port}/{database}"


def _resolve_pg_url() -> str:
    url = os.getenv("DATABASE_URL")
    if not url:
        _build_pg_url_from_parts()
        url = os.getenv("DATABASE_URL")
    if not url:
        raise RuntimeError(
            "DATABASE_URL not set. Set DATABASE_URL or PG_USER/PG_PASSWORD/PG_HOST/PG_PORT/PG_DATABASE."
        )
    return url


def _resolve_doris_url() -> str:
    for key in ("DORIS_URL", "DORIS_DSN"):
        v = os.getenv(key)
        if v:
            return v
    user = os.getenv("DORIS_USER", "root")
    password = os.getenv("DORIS_PASSWORD", "")
    host = os.getenv("DORIS_HOST", "127.0.0.1")
    port = os.getenv("DORIS_PORT", "9030")
    database = os.getenv("DORIS_DATABASE", os.getenv("DORIS_DB", "dream"))
    auth = f"{user}:{password}@" if password else f"{user}@"
    return f"mysql+aiomysql://{auth}{host}:{port}/{database}"


# ---------------------------------------------------------------------------
# Migration resolution helpers
# ---------------------------------------------------------------------------


def _resolve_pg_migrations(version_file: str | None) -> list:
    from repo_core.migrations.registry import get_postgres_migrations, get_postgres_migrations_up_to
    if version_file:
        return get_postgres_migrations_up_to(version_file)
    return get_postgres_migrations()


def _resolve_doris_migrations(version_file: str | None) -> list:
    from repo_core.migrations.registry import get_doris_migrations, get_doris_migrations_up_to
    if version_file:
        return get_doris_migrations_up_to(version_file)
    return get_doris_migrations()


# ---------------------------------------------------------------------------
# PostgreSQL commands
# ---------------------------------------------------------------------------


async def _pg_upgrade(version_file: str | None = None) -> None:
    from db_core import create_postgres_client, ensure_pg_database_exists
    from repo_core.migrations.runner import MigrationRunner

    url = _resolve_pg_url()
    await ensure_pg_database_exists(url)
    client = create_postgres_client(url=url)
    runner = MigrationRunner(client, _resolve_pg_migrations(version_file), dialect="postgresql")
    await client.connect()
    try:
        await runner.ensure_migrations_table()
        await runner.upgrade()
    finally:
        await client.disconnect()


async def _pg_downgrade(target: str | None, version_file: str | None = None) -> None:
    from db_core import create_postgres_client, ensure_pg_database_exists
    from repo_core.migrations.runner import MigrationRunner

    url = _resolve_pg_url()
    await ensure_pg_database_exists(url)
    client = create_postgres_client(url=url)
    runner = MigrationRunner(client, _resolve_pg_migrations(version_file), dialect="postgresql")
    await client.connect()
    try:
        await runner.ensure_migrations_table()
        await runner.downgrade(target)
    finally:
        await client.disconnect()


async def _pg_status() -> None:
    from db_core import create_postgres_client, ensure_pg_database_exists
    from repo_core.migrations.registry import get_postgres_migrations
    from repo_core.migrations.runner import MigrationRunner

    url = _resolve_pg_url()
    await ensure_pg_database_exists(url)
    client = create_postgres_client(url=url)
    runner = MigrationRunner(client, get_postgres_migrations(), dialect="postgresql")
    await client.connect()
    try:
        await runner.ensure_migrations_table()
        rows = await runner.status()
    finally:
        await client.disconnect()

    _print_status(rows)


# ---------------------------------------------------------------------------
# Doris commands
# ---------------------------------------------------------------------------


async def _doris_upgrade(version_file: str | None = None) -> None:
    from db_core import create_doris_client, ensure_doris_database_exists
    from repo_core.migrations.runner import MigrationRunner

    url = _resolve_doris_url()
    await ensure_doris_database_exists(url)
    client = create_doris_client(url=url)
    runner = MigrationRunner(client, _resolve_doris_migrations(version_file), dialect="doris")
    await client.connect()
    try:
        await runner.ensure_migrations_table()
        await runner.upgrade()
    finally:
        await client.disconnect()


async def _doris_downgrade(target: str | None, version_file: str | None = None) -> None:
    from db_core import create_doris_client, ensure_doris_database_exists
    from repo_core.migrations.runner import MigrationRunner

    url = _resolve_doris_url()
    await ensure_doris_database_exists(url)
    client = create_doris_client(url=url)
    runner = MigrationRunner(client, _resolve_doris_migrations(version_file), dialect="doris")
    await client.connect()
    try:
        await runner.ensure_migrations_table()
        await runner.downgrade(target)
    finally:
        await client.disconnect()


async def _doris_status() -> None:
    from db_core import create_doris_client, ensure_doris_database_exists
    from repo_core.migrations.registry import get_doris_migrations
    from repo_core.migrations.runner import MigrationRunner

    url = _resolve_doris_url()
    await ensure_doris_database_exists(url)
    client = create_doris_client(url=url)
    runner = MigrationRunner(client, get_doris_migrations(), dialect="doris")
    await client.connect()
    try:
        await runner.ensure_migrations_table()
        rows = await runner.status()
    finally:
        await client.disconnect()

    _print_status(rows)


# ---------------------------------------------------------------------------
# Generate command
# ---------------------------------------------------------------------------


def _run_generate(message: str, dialect: str = "postgres") -> None:
    """Scaffold a new migration version file.

    Args:
        message: Human-readable description for the migration.
        dialect: ``"postgres"`` (default) or ``"doris"``.
    """
    subdir = "doris" if dialect == "doris" else "postgres"
    versions_dir = Path(__file__).resolve().parent / "versions" / subdir
    versions_dir.mkdir(parents=True, exist_ok=True)

    # Find next version number within the dialect subdirectory
    existing = sorted(versions_dir.glob("[0-9]*_*.py"))
    if existing:
        last_prefix = existing[-1].stem.split("_", 1)[0]
        try:
            next_num = int(last_prefix) + 1
        except ValueError:
            next_num = 1
    else:
        next_num = 1

    slug = message.lower().replace(" ", "_").replace("-", "_")[:60]
    filename = f"{next_num:04d}_{slug}.py"
    filepath = versions_dir / filename

    if subdir == "doris":
        template = f'''"""{message} — migration {next_num:04d}."""

from repo_core.migrations.builder import MigrationBuilder
from repo_core.migrations.runner import Migration

VERSION = "{next_num:04d}"


def get_doris_migrations() -> list[Migration]:
    return [
        # Migration(
        #     version="{slug}_0001",
        #     description="{message}",
        #     dialect="doris",
        #     up=(
        #         MigrationBuilder("table_name", dialect="doris")
        #         .bigint("id", nullable=False)
        #         .text("name", nullable=False)
        #         .doris_unique_key(["id"])
        #         .doris_distributed_by("id", buckets=10)
        #         .create_table()
        #     ),
        #     down=MigrationBuilder("table_name", dialect="doris").drop_table(),
        # ),
    ]
'''
    else:
        template = f'''"""{message} — migration {next_num:04d}."""

from repo_core.migrations.builder import MigrationBuilder
from repo_core.migrations.runner import Migration

VERSION = "{next_num:04d}"


def get_postgres_migrations() -> list[Migration]:
    return [
        # Migration(
        #     version="{slug}_0001",
        #     description="{message}",
        #     up=(
        #         MigrationBuilder("table_name")
        #         .uuid("id", primary_key=True, nullable=False)
        #         .text("name", nullable=False)
        #         .create_table()
        #     ),
        #     down=MigrationBuilder("table_name").drop_table(),
        # ),
    ]
'''

    filepath.write_text(template)
    print(f"  Created {filepath}")


# ---------------------------------------------------------------------------
# Output helpers
# ---------------------------------------------------------------------------


def _print_status(rows: list[dict]) -> None:
    col_v = max(len(r["version"]) for r in rows) if rows else 7
    col_d = max(len(r["description"]) for r in rows) if rows else 11
    col_v = max(col_v, 7)
    col_d = max(col_d, 11)
    header = f"{'Version':<{col_v}}  {'Status':<9}  {'Rollback':<8}  Description"
    print(f"\n{header}")
    print("-" * (len(header) + 10))
    for r in rows:
        rb = "yes" if r["has_rollback"] else "no"
        status_str = r["status"]
        if status_str == "applied":
            marker = "✓ applied"
        else:
            marker = "  pending"
        print(f"{r['version']:<{col_v}}  {marker:<9}  {rb:<8}  {r['description']}")
    print()


# ---------------------------------------------------------------------------
# Argument parser
# ---------------------------------------------------------------------------


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="python -m repo_core.migrations",
        description="repo_core schema migration runner.",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # PostgreSQL
    pg_up = sub.add_parser("upgrade", help="Apply pending PostgreSQL migrations")
    pg_up.add_argument(
        "--file", dest="version_file", default=None,
        help="Apply migrations only up to this version file (e.g., '0002')",
    )

    pg_down = sub.add_parser("downgrade", help="Roll back PostgreSQL migrations")
    pg_down.add_argument(
        "target", nargs="?", default=None,
        help="Roll back to (but not including) this version. Omit to roll back one step.",
    )
    pg_down.add_argument(
        "--file", dest="version_file", default=None,
        help="Limit to migrations from version files up to this prefix",
    )

    sub.add_parser("status", help="Show PostgreSQL migration status")

    # Generate
    gen = sub.add_parser("generate", help="Scaffold a new migration version file")
    gen.add_argument("-m", "--message", required=True, help="Description for the new migration")
    gen.add_argument(
        "-d", "--dialect", default="postgres",
        choices=["postgres", "doris"],
        help="Dialect for the new migration: 'postgres' (default) or 'doris'",
    )

    # Doris
    doris_up = sub.add_parser("doris-upgrade", help="Apply pending Doris warehouse migrations")
    doris_up.add_argument(
        "--file", dest="version_file", default=None,
        help="Apply migrations only up to this version file",
    )

    doris_down = sub.add_parser("doris-downgrade", help="Roll back Doris warehouse migrations")
    doris_down.add_argument("target", nargs="?", default=None)
    doris_down.add_argument("--file", dest="version_file", default=None)

    sub.add_parser("doris-status", help="Show Doris warehouse migration status")

    return parser


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    _load_env()
    parser = _build_parser()
    args = parser.parse_args()

    vf = getattr(args, "version_file", None)

    try:
        if args.command == "upgrade":
            asyncio.run(_pg_upgrade(vf))
        elif args.command == "downgrade":
            asyncio.run(_pg_downgrade(getattr(args, "target", None), vf))
        elif args.command == "status":
            asyncio.run(_pg_status())
        elif args.command == "generate":
            _run_generate(args.message, dialect=args.dialect)
        elif args.command == "doris-upgrade":
            asyncio.run(_doris_upgrade(vf))
        elif args.command == "doris-downgrade":
            asyncio.run(_doris_downgrade(getattr(args, "target", None), vf))
        elif args.command == "doris-status":
            asyncio.run(_doris_status())
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
