"""
UUID Migration Script

Migrates all research_id fields from Integer to String (UUID) format.
This creates a more consistent and scalable ID system across the application.
"""

import sqlite3
import uuid
from pathlib import Path
from loguru import logger


def get_database_path():
    """Get the path to the SQLite database."""
    data_dir = Path(__file__).parents[3] / "data"
    return data_dir / "ldr.db"


def migrate_to_uuid():
    """
    Migrate all research_id fields from integers to UUIDs.

    Strategy:
    1. Add new UUID columns alongside existing integer columns
    2. Generate UUIDs for existing data (or keep as string versions of integers)
    3. Update foreign key relationships
    4. Drop old integer columns and rename UUID columns
    """
    db_path = get_database_path()

    if not db_path.exists():
        logger.info("Database doesn't exist yet, migration not needed")
        return

    logger.info(f"Starting UUID migration on {db_path}")

    conn = sqlite3.connect(db_path)
    conn.execute(
        "PRAGMA foreign_keys = OFF"
    )  # Disable FK constraints during migration

    try:
        cursor = conn.cursor()

        # 1. Migrate research_history table (main research IDs)
        logger.info("Migrating research_history table...")

        # Check if the table exists and has the old structure
        cursor.execute("PRAGMA table_info(research_history)")
        columns = cursor.fetchall()
        has_uuid_id = any(col[1] == "uuid_id" for col in columns)

        if not has_uuid_id:
            # Add UUID column
            cursor.execute(
                "ALTER TABLE research_history ADD COLUMN uuid_id TEXT"
            )

            # Generate UUIDs for existing records (convert integer ID to UUID format)
            cursor.execute("SELECT id FROM research_history")
            existing_ids = cursor.fetchall()

            for (old_id,) in existing_ids:
                # Generate a deterministic UUID based on the old ID
                new_uuid = str(
                    uuid.uuid5(uuid.NAMESPACE_OID, f"research_{old_id}")
                )
                cursor.execute(
                    "UPDATE research_history SET uuid_id = ? WHERE id = ?",
                    (new_uuid, old_id),
                )

            logger.info(
                f"Generated UUIDs for {len(existing_ids)} research records"
            )

        # 2. Migrate metrics tables
        logger.info("Migrating metrics tables...")

        # Token usage table
        cursor.execute("PRAGMA table_info(token_usage)")
        columns = cursor.fetchall()
        has_uuid_research_id = any(
            col[1] == "uuid_research_id" for col in columns
        )

        if not has_uuid_research_id:
            cursor.execute(
                "ALTER TABLE token_usage ADD COLUMN uuid_research_id TEXT"
            )

            # Convert existing research_ids to UUIDs (deterministic conversion)
            cursor.execute(
                "SELECT DISTINCT research_id FROM token_usage WHERE research_id IS NOT NULL"
            )
            research_ids = cursor.fetchall()

            for (research_id,) in research_ids:
                if research_id:
                    new_uuid = str(
                        uuid.uuid5(
                            uuid.NAMESPACE_OID, f"research_{research_id}"
                        )
                    )
                    cursor.execute(
                        "UPDATE token_usage SET uuid_research_id = ? WHERE research_id = ?",
                        (new_uuid, research_id),
                    )

            logger.info(
                f"Migrated {len(research_ids)} research IDs in token_usage"
            )

        # Model usage table
        cursor.execute("PRAGMA table_info(model_usage)")
        columns = cursor.fetchall()
        has_uuid_research_id = any(
            col[1] == "uuid_research_id" for col in columns
        )

        if not has_uuid_research_id:
            cursor.execute(
                "ALTER TABLE model_usage ADD COLUMN uuid_research_id TEXT"
            )

            cursor.execute(
                "SELECT DISTINCT research_id FROM model_usage WHERE research_id IS NOT NULL"
            )
            research_ids = cursor.fetchall()

            for (research_id,) in research_ids:
                if research_id:
                    new_uuid = str(
                        uuid.uuid5(
                            uuid.NAMESPACE_OID, f"research_{research_id}"
                        )
                    )
                    cursor.execute(
                        "UPDATE model_usage SET uuid_research_id = ? WHERE research_id = ?",
                        (new_uuid, research_id),
                    )

            logger.info(
                f"Migrated {len(research_ids)} research IDs in model_usage"
            )

        # Search calls table
        cursor.execute("PRAGMA table_info(search_calls)")
        columns = cursor.fetchall()
        has_uuid_research_id = any(
            col[1] == "uuid_research_id" for col in columns
        )

        if not has_uuid_research_id:
            cursor.execute(
                "ALTER TABLE search_calls ADD COLUMN uuid_research_id TEXT"
            )

            cursor.execute(
                "SELECT DISTINCT research_id FROM search_calls WHERE research_id IS NOT NULL"
            )
            research_ids = cursor.fetchall()

            for (research_id,) in research_ids:
                if research_id:
                    new_uuid = str(
                        uuid.uuid5(
                            uuid.NAMESPACE_OID, f"research_{research_id}"
                        )
                    )
                    cursor.execute(
                        "UPDATE search_calls SET uuid_research_id = ? WHERE research_id = ?",
                        (new_uuid, research_id),
                    )

            logger.info(
                f"Migrated {len(research_ids)} research IDs in search_calls"
            )

        # 3. Migrate benchmark tables
        logger.info("Migrating benchmark tables...")

        # Check if benchmark_results table exists
        cursor.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name='benchmark_results'"
        )
        if cursor.fetchone():
            cursor.execute("PRAGMA table_info(benchmark_results)")
            columns = cursor.fetchall()
            has_uuid_research_id = any(
                col[1] == "uuid_research_id" for col in columns
            )

            if not has_uuid_research_id:
                cursor.execute(
                    "ALTER TABLE benchmark_results ADD COLUMN uuid_research_id TEXT"
                )

                cursor.execute(
                    "SELECT DISTINCT research_id FROM benchmark_results WHERE research_id IS NOT NULL"
                )
                research_ids = cursor.fetchall()

                for (research_id,) in research_ids:
                    if research_id:
                        new_uuid = str(
                            uuid.uuid5(
                                uuid.NAMESPACE_OID, f"research_{research_id}"
                            )
                        )
                        cursor.execute(
                            "UPDATE benchmark_results SET uuid_research_id = ? WHERE research_id = ?",
                            (new_uuid, research_id),
                        )

                logger.info(
                    f"Migrated {len(research_ids)} research IDs in benchmark_results"
                )

        # Commit all changes
        conn.commit()
        logger.info("UUID migration completed successfully!")

        # Note: We're keeping both old and new columns for now
        # The application will use the new UUID columns
        # Old columns can be dropped in a future migration once everything is stable

    except Exception as e:
        logger.error(f"Error during UUID migration: {e}")
        conn.rollback()
        raise
    finally:
        conn.execute("PRAGMA foreign_keys = ON")  # Re-enable FK constraints
        conn.close()


def cleanup_old_columns():
    """
    Cleanup migration - drops old integer columns after UUID migration is stable.
    Run this only after confirming the UUID migration is working correctly.
    """
    logger.warning(
        "This will permanently remove old integer research_id columns!"
    )
    logger.warning("Make sure to backup your database before running this!")

    db_path = get_database_path()
    conn = sqlite3.connect(db_path)
    conn.execute("PRAGMA foreign_keys = OFF")

    try:
        # For SQLite, we need to recreate tables to drop columns
        # This is complex, so we'll leave old columns for now
        # They can be cleaned up manually if needed

        logger.info("Cleanup deferred - old columns remain for safety")

    finally:
        conn.execute("PRAGMA foreign_keys = ON")
        conn.close()


if __name__ == "__main__":
    migrate_to_uuid()
