"""
Convert research_id columns from Integer to String.

This migration converts existing integer research_id values to string format
while preserving all existing data. New records will use UUID strings.
"""

import sqlite3
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 convert_research_id_to_string():
    """
    Convert research_id columns from Integer to String in all tables.
    Preserves existing data by converting integer IDs to string format.
    """
    db_path = get_database_path()

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

    logger.info(f"Converting research_id columns to string in {db_path}")

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

    try:
        cursor = conn.cursor()

        # List of tables that have research_id columns
        tables_to_migrate = [
            "token_usage",
            "model_usage",
            "search_calls",
            "benchmark_results",  # If it exists
        ]

        for table_name in tables_to_migrate:
            logger.info(f"Converting {table_name} table...")

            # Check if table exists
            cursor.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
                (table_name,),
            )
            if not cursor.fetchone():
                logger.info(f"Table {table_name} does not exist, skipping")
                continue

            # Check if research_id column exists
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns = cursor.fetchall()
            has_research_id = any(col[1] == "research_id" for col in columns)

            if not has_research_id:
                logger.info(
                    f"Table {table_name} does not have research_id column, skipping"
                )
                continue

            # For SQLite, we need to recreate the table to change column type
            # 1. Create new table with string research_id
            # 2. Copy data with research_id converted to string
            # 3. Drop old table and rename new table

            # Get the current table schema
            cursor.execute(
                f"SELECT sql FROM sqlite_master WHERE type='table' AND name='{table_name}'"
            )
            create_sql = cursor.fetchone()[0]

            # Create new table name
            new_table_name = f"{table_name}_new"

            # Modify the CREATE TABLE statement to change research_id to TEXT
            new_create_sql = create_sql.replace(
                f"CREATE TABLE {table_name}", f"CREATE TABLE {new_table_name}"
            )
            new_create_sql = new_create_sql.replace(
                "research_id INTEGER", "research_id TEXT"
            )
            new_create_sql = new_create_sql.replace(
                "research_id INT", "research_id TEXT"
            )

            # Create the new table
            cursor.execute(new_create_sql)

            # Copy data from old table to new table, converting research_id to string
            cursor.execute(f"SELECT * FROM {table_name}")
            old_rows = cursor.fetchall()

            if old_rows:
                # Get column names
                cursor.execute(f"PRAGMA table_info({table_name})")
                columns = cursor.fetchall()
                column_names = [col[1] for col in columns]
                research_id_index = (
                    column_names.index("research_id")
                    if "research_id" in column_names
                    else -1
                )

                # Prepare insert statement
                placeholders = ",".join(["?" for _ in column_names])
                insert_sql = f"INSERT INTO {new_table_name} ({','.join(column_names)}) VALUES ({placeholders})"

                # Convert rows and insert
                converted_rows = []
                for row in old_rows:
                    row_list = list(row)
                    # Convert research_id to string if it's not None
                    if (
                        research_id_index >= 0
                        and row_list[research_id_index] is not None
                    ):
                        row_list[research_id_index] = str(
                            row_list[research_id_index]
                        )
                    converted_rows.append(tuple(row_list))

                cursor.executemany(insert_sql, converted_rows)
                logger.info(
                    f"Converted {len(converted_rows)} rows in {table_name}"
                )

            # Drop old table and rename new table
            cursor.execute(f"DROP TABLE {table_name}")
            cursor.execute(
                f"ALTER TABLE {new_table_name} RENAME TO {table_name}"
            )

            logger.info(
                f"Successfully converted {table_name} research_id to string"
            )

        # Commit all changes
        conn.commit()
        logger.info("All research_id columns converted to string successfully!")

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


if __name__ == "__main__":
    convert_research_id_to_string()
