#! /usr/bin/python3
"""Migration runner for the garments schema.

Applies schema/garments.sql (base schema, fresh installs only) and
schema/migrations/NNN_*.sql files in lexicographic order, skipping any
file already recorded in garments.schema_migrations.

Usage:
    scripts/migrate [options]

Options:
    --config PATH       Garments config file.
                        Defaults to GARMENTS_CONFIG env var or /etc/garments.
    --schema-dir PATH   Path to the schema/ directory.
                        Defaults to schema/ relative to the repository root.
    --dry-run           Print what would be applied without making any changes.
    --backfill          Mark all migration files as applied without running them.
                        Use once on an existing installation to adopt the runner.
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path
from typing import Optional

import psycopg2
import psycopg2.extras

# Locate the repository root relative to this script (scripts/ → repo root).
_REPO_ROOT = Path(__file__).resolve().parent.parent
_DEFAULT_SCHEMA_DIR = _REPO_ROOT / 'schema'

# Ensure the garments package is importable when the script is run directly
# (i.e. not via 'python3 -m'), which does not add the repo root to sys.path.
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

_CREATE_MIGRATIONS_TABLE = """
CREATE TABLE IF NOT EXISTS garments.schema_migrations (
    filename   TEXT        PRIMARY KEY,
    applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
);

COMMENT ON TABLE garments.schema_migrations IS
    'Tracks which migration SQL files have been applied. '
    'Managed by scripts/migrate — do not edit manually.';
"""


# ============================================================
# ARGUMENT PARSING
# ============================================================

def build_parser() -> argparse.ArgumentParser:
    """Return the argument parser for the migration runner."""
    parser = argparse.ArgumentParser(
        prog='migrate',
        description='Apply pending garments schema migrations.',
    )
    parser.add_argument(
        '--config',
        default=None,
        metavar='PATH',
        help=(
            'Garments config file. '
            'Defaults to GARMENTS_CONFIG env var or /etc/garments.'
        ),
    )
    parser.add_argument(
        '--schema-dir',
        default=str(_DEFAULT_SCHEMA_DIR),
        metavar='PATH',
        help='Path to the schema/ directory (default: %(default)s).',
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        default=False,
        help='Print pending migrations without applying them.',
    )
    parser.add_argument(
        '--backfill',
        action='store_true',
        default=False,
        help=(
            'Mark all migration files as applied without running them. '
            'Use once on an existing installation.'
        ),
    )
    return parser


# ============================================================
# PURE HELPERS  (no I/O — easily unit-tested)
# ============================================================

def find_migration_files(migrations_dir: Path) -> list[Path]:
    """Return .sql files in migrations_dir sorted lexicographically."""
    return sorted(migrations_dir.glob('*.sql'))


def find_pending(
    all_files: list[Path],
    applied: set[str],
) -> list[Path]:
    """Return files whose names are not in the applied set."""
    return [f for f in all_files if f.name not in applied]


# ============================================================
# DATABASE HELPERS
# ============================================================

def _connect(dsn: dict) -> psycopg2.extensions.connection:
    conn = psycopg2.connect(**dsn)
    conn.cursor_factory = psycopg2.extras.RealDictCursor
    return conn


def schema_exists(conn: psycopg2.extensions.connection) -> bool:
    """Return True if the garments schema is present in the database."""
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT 1 FROM information_schema.schemata
            WHERE schema_name = 'garments'
            """,
        )
        return cur.fetchone() is not None


def ensure_migrations_table(conn: psycopg2.extensions.connection) -> None:
    """Create garments.schema_migrations if it does not already exist."""
    with conn.cursor() as cur:
        cur.execute(_CREATE_MIGRATIONS_TABLE)
    conn.commit()


def get_applied(conn: psycopg2.extensions.connection) -> set[str]:
    """Return the set of filenames already recorded in schema_migrations."""
    with conn.cursor() as cur:
        cur.execute('SELECT filename FROM garments.schema_migrations')
        return {row['filename'] for row in cur.fetchall()}


def record_migration(
    conn: psycopg2.extensions.connection,
    filename: str,
) -> None:
    """Insert a filename into schema_migrations to mark it as applied."""
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO garments.schema_migrations (filename) VALUES (%s)',
            (filename,),
        )
    conn.commit()


# ============================================================
# PSQL SUBPROCESS
# ============================================================

def _build_psql_env(dsn: dict) -> dict:
    """Return os.environ with PGPASSWORD injected if a password is set."""
    env = os.environ.copy()
    if dsn.get('password'):
        env['PGPASSWORD'] = dsn['password']
    return env


def run_psql(dsn: dict, sql_file: Path) -> None:
    """Apply a SQL file via psql, raising RuntimeError on failure.

    Uses the same connection parameters as the psycopg2 connection so
    the two are always in sync.
    """
    cmd = [
        'psql',
        '--host',     dsn.get('host', 'localhost'),
        '--port',     str(dsn.get('port', 5432)),
        '--username', dsn['user'],
        '--dbname',   dsn['dbname'],
        '--no-password',
        '--file',     str(sql_file),
    ]
    result = subprocess.run(
        cmd,
        env=_build_psql_env(dsn),
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(
            f'psql exited {result.returncode} applying {sql_file.name}:\n'
            f'{result.stderr.strip()}'
        )
    if result.stdout.strip():
        print(result.stdout.strip())


# ============================================================
# MAIN
# ============================================================

def _load_dsn(config_path: Optional[str]) -> dict:
    """Load DSN kwargs from the garments config."""
    # Honour explicit --config, then GARMENTS_CONFIG, then /etc/garments.
    if config_path:
        os.environ['GARMENTS_CONFIG'] = config_path

    # Import here so the env var override is in place first.
    from garments.config import db_dsn
    return db_dsn()


def run(argv: Optional[list[str]] = None) -> int:
    """Entry point. Returns an exit code."""
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.dry_run and args.backfill:
        print('error: --dry-run and --backfill are mutually exclusive.',
              file=sys.stderr)
        return 2

    schema_dir = Path(args.schema_dir)
    base_sql = schema_dir / 'garments.sql'
    migrations_dir = schema_dir / 'migrations'

    if not base_sql.exists():
        print(f'error: base schema not found: {base_sql}', file=sys.stderr)
        return 2
    if not migrations_dir.is_dir():
        print(f'error: migrations directory not found: {migrations_dir}',
              file=sys.stderr)
        return 2

    dsn = _load_dsn(args.config)
    conn = _connect(dsn)

    # ---- Bootstrap: apply base schema on fresh databases ----
    if not schema_exists(conn):
        if args.dry_run:
            print(f'[dry-run] Would apply base schema: {base_sql.name}')
        else:
            print(f'Applying base schema: {base_sql.name} ...')
            try:
                run_psql(dsn, base_sql)
            except RuntimeError as exc:
                print(f'error: {exc}', file=sys.stderr)
                conn.close()
                return 1
            print('  OK')
            # Reconnect now that the schema exists.
            conn.close()
            conn = _connect(dsn)

    # ---- Ensure tracking table exists ----
    if not args.dry_run:
        ensure_migrations_table(conn)

    # ---- Determine pending migrations ----
    all_files = find_migration_files(migrations_dir)

    if args.dry_run:
        # In dry-run we cannot query schema_migrations (table may not exist
        # on a truly fresh DB that hasn't had garments.sql applied yet).
        # Show all files as pending in that case.
        try:
            applied = get_applied(conn)
        except Exception:
            applied = set()
        pending = find_pending(all_files, applied)
        if not pending:
            print('No pending migrations.')
        else:
            print(f'{len(pending)} pending migration(s):')
            for f in pending:
                print(f'  {f.name}')
        conn.close()
        return 0

    applied = get_applied(conn)
    pending = find_pending(all_files, applied)

    if not pending:
        print('No pending migrations.')
        conn.close()
        return 0

    print(
        f'{"Backfilling" if args.backfill else "Applying"} '
        f'{len(pending)} pending migration(s):'
    )

    for migration in pending:
        if args.backfill:
            print(f'  [backfill] {migration.name}')
            record_migration(conn, migration.name)
        else:
            print(f'  Applying {migration.name} ...', end=' ', flush=True)
            try:
                run_psql(dsn, migration)
            except RuntimeError as exc:
                print()
                print(f'error: {exc}', file=sys.stderr)
                conn.close()
                return 1
            record_migration(conn, migration.name)
            print('OK')

    conn.close()
    return 0


def main() -> None:
    """CLI entry point."""
    sys.exit(run())


if __name__ == '__main__':
    main()
