#!/usr/bin/env python3
"""Fix PostgreSQL sequences and migrate datetime columns after SQLite migration.

This script performs two operations:
1. Migrates DateTime columns to TIMESTAMP WITH TIME ZONE for timezone-aware support
2. Resets all ID sequence values to match the current maximum ID in each table

The sequence reset is necessary because the migration script inserts records
with explicit IDs, which doesn't automatically advance PostgreSQL sequences.

The datetime migration ensures proper handling of timezone-aware datetimes,
which is required by the application code.

Schema Support:
    The script automatically detects and uses the schema configured in your config file.
    If you have database.schema: "alpinebits" in your config, it will work with that schema.

Usage:
    # Using default config.yaml (includes schema if configured)
    uv run python -m alpine_bits_python.util.fix_postgres_sequences

    # Using a specific config file (with schema support)
    uv run python -m alpine_bits_python.util.fix_postgres_sequences \
        --config config/postgres.yaml

    # Using DATABASE_URL environment variable (schema from config or DATABASE_SCHEMA env var)
    DATABASE_URL="postgresql+asyncpg://user:pass@host/db" \
    DATABASE_SCHEMA="alpinebits" \
        uv run python -m alpine_bits_python.util.fix_postgres_sequences

    # Using command line argument (schema from config)
    uv run python -m alpine_bits_python.util.fix_postgres_sequences \
        --database-url postgresql+asyncpg://user:pass@host/db
"""

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

# Add parent directory to path so we can import alpine_bits_python
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

import yaml
from sqlalchemy import text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from alpine_bits_python.db import get_database_schema, get_database_url
from alpine_bits_python.logging_config import get_logger

_LOGGER = get_logger(__name__)


async def migrate_datetime_columns(session, schema_prefix: str = "") -> None:
    """Migrate DateTime columns to TIMESTAMP WITH TIME ZONE.

    This updates the columns to properly handle timezone-aware datetimes.

    Args:
        session: Database session
        schema_prefix: Schema prefix (e.g., "alpinebits." or "")

    """
    _LOGGER.info("\nMigrating DateTime columns to timezone-aware...")

    datetime_columns = [
        ("hashed_customers", "created_at"),
        ("reservations", "created_at"),
        ("acked_requests", "timestamp"),
    ]

    for table_name, column_name in datetime_columns:
        full_table = f"{schema_prefix}{table_name}"
        _LOGGER.info(f"  {full_table}.{column_name}: Converting to TIMESTAMPTZ")
        await session.execute(
            text(
                f"ALTER TABLE {full_table} "
                f"ALTER COLUMN {column_name} TYPE TIMESTAMP WITH TIME ZONE"
            )
        )

    await session.commit()
    _LOGGER.info("✓ DateTime columns migrated to timezone-aware")


async def fix_sequences(database_url: str, schema_name: str = None) -> None:
    """Fix PostgreSQL sequences to match current max IDs and migrate datetime columns.

    Args:
        database_url: PostgreSQL database URL
        schema_name: Schema name (e.g., "alpinebits") or None for public

    """
    _LOGGER.info("=" * 70)
    _LOGGER.info("PostgreSQL Migration & Sequence Fix")
    _LOGGER.info("=" * 70)
    _LOGGER.info(
        "Database: %s",
        database_url.split("@")[-1] if "@" in database_url else database_url,
    )
    if schema_name:
        _LOGGER.info("Schema: %s", schema_name)
    _LOGGER.info("=" * 70)

    # Create engine and session with schema support
    connect_args = {}
    if schema_name:
        connect_args = {
            "server_settings": {"search_path": f"{schema_name},public"}
        }

    engine = create_async_engine(database_url, echo=False, connect_args=connect_args)
    SessionMaker = async_sessionmaker(engine, expire_on_commit=False)

    # Determine schema prefix for SQL statements
    schema_prefix = f"{schema_name}." if schema_name else ""

    try:
        # Migrate datetime columns first
        async with SessionMaker() as session:
            await migrate_datetime_columns(session, schema_prefix)

        # Then fix sequences
        async with SessionMaker() as session:
            # List of tables and their sequence names
            tables = [
                ("customers", "customers_id_seq"),
                ("hashed_customers", "hashed_customers_id_seq"),
                ("reservations", "reservations_id_seq"),
                ("acked_requests", "acked_requests_id_seq"),
                ("conversions", "conversions_id_seq"),
            ]

            _LOGGER.info("\nResetting sequences...")
            for table_name, sequence_name in tables:
                full_table = f"{schema_prefix}{table_name}"
                full_sequence = f"{schema_prefix}{sequence_name}"

                # Get current max ID
                result = await session.execute(
                    text(f"SELECT MAX(id) FROM {full_table}")
                )
                max_id = result.scalar()

                # Get current sequence value
                result = await session.execute(
                    text(f"SELECT last_value FROM {full_sequence}")
                )
                current_seq = result.scalar()

                if max_id is None:
                    _LOGGER.info(f"  {full_table}: empty table, setting sequence to 1")
                    await session.execute(
                        text(f"SELECT setval('{full_sequence}', 1, false)")
                    )
                elif current_seq <= max_id:
                    new_seq = max_id + 1
                    _LOGGER.info(
                        f"  {full_table}: max_id={max_id}, "
                        f"old_seq={current_seq}, new_seq={new_seq}"
                    )
                    await session.execute(
                        text(f"SELECT setval('{full_sequence}', {new_seq}, false)")
                    )
                else:
                    _LOGGER.info(
                        f"  {full_table}: sequence already correct "
                        f"(max_id={max_id}, seq={current_seq})"
                    )

            await session.commit()

        _LOGGER.info("\n" + "=" * 70)
        _LOGGER.info("✓ Migration completed successfully!")
        _LOGGER.info("=" * 70)
        _LOGGER.info("\nChanges applied:")
        _LOGGER.info("  1. DateTime columns are now timezone-aware (TIMESTAMPTZ)")
        _LOGGER.info("  2. Sequences are reset to match current max IDs")
        _LOGGER.info("\nYou can now insert new records without conflicts.")

    except Exception as e:
        _LOGGER.exception("Failed to fix sequences: %s", e)
        raise

    finally:
        await engine.dispose()


async def main():
    """Run the sequence fix."""
    parser = argparse.ArgumentParser(
        description="Fix PostgreSQL sequences after SQLite migration",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "--database-url",
        help="PostgreSQL database URL (default: from config or DATABASE_URL env var)",
    )
    parser.add_argument(
        "--config",
        help=(
            "Path to config file containing PostgreSQL database URL "
            "(keeps password out of bash history)"
        ),
    )

    args = parser.parse_args()

    # Determine database URL and schema
    schema_name = None
    if args.config:
        # Load config file as plain YAML (no validation)
        _LOGGER.info("Loading database config from: %s", args.config)
        try:
            with open(args.config) as f:
                config = yaml.safe_load(f)

            # Get database URL
            if not config or "database" not in config or "url" not in config["database"]:
                _LOGGER.error("Config file must contain database.url")
                sys.exit(1)

            database_url = config["database"]["url"]
            schema_name = config.get("database", {}).get("schema")
            _LOGGER.info("Successfully loaded config")
        except FileNotFoundError:
            _LOGGER.error("Config file not found: %s", args.config)
            sys.exit(1)
        except yaml.YAMLError as e:
            _LOGGER.error("Failed to parse YAML config: %s", e)
            sys.exit(1)
    elif args.database_url:
        database_url = args.database_url
        # Get schema from environment variable
        schema_name = os.environ.get("DATABASE_SCHEMA")
    else:
        # Try environment variable or default config.yaml
        database_url = os.environ.get("DATABASE_URL")
        schema_name = os.environ.get("DATABASE_SCHEMA")

        if not database_url:
            # Try to load from default config.yaml as plain YAML
            try:
                with open("config/config.yaml") as f:
                    config = yaml.safe_load(f)
                    database_url = config.get("database", {}).get("url")
                    if not schema_name:
                        schema_name = config.get("database", {}).get("schema")
            except Exception:
                pass  # Ignore if default config doesn't exist

        if not database_url:
            _LOGGER.error("No database URL provided")
            _LOGGER.error("Provide via --config, --database-url, or DATABASE_URL env var")
            sys.exit(1)

    if "postgresql" not in database_url and "postgres" not in database_url:
        _LOGGER.error("This script only works with PostgreSQL databases.")
        url_type = database_url.split("+")[0] if "+" in database_url else "unknown"
        _LOGGER.error("Current database URL type detected: %s", url_type)
        _LOGGER.error("\nSpecify PostgreSQL database using one of:")
        _LOGGER.error("  - --config config/postgres.yaml")
        _LOGGER.error("  - DATABASE_URL environment variable")
        _LOGGER.error("  - --database-url postgresql+asyncpg://user:pass@host/db")
        sys.exit(1)

    # Run the fix
    await fix_sequences(database_url, schema_name)


if __name__ == "__main__":
    asyncio.run(main())
