"""
Schema upgrade script for Local Deep Research database.
Handles schema upgrades for existing ldr.db databases.
"""

import os
import sys

from loguru import logger
from sqlalchemy import create_engine, inspect, text

# Add the parent directory to sys.path to allow relative imports
sys.path.append(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
)

try:
    from src.local_deep_research.web.models.database import DB_PATH
except ImportError:
    # Fallback path if import fails
    current_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.abspath(
        os.path.join(current_dir, "..", "..", "..", "..")
    )
    DB_PATH = os.path.join(project_root, "src", "data", "ldr.db")

from .models import Base, ResearchStrategy, RateLimitAttempt, RateLimitEstimate


def remove_research_log_table(engine):
    """
    Remove the redundant research_log table if it exists

    Args:
        engine: SQLAlchemy engine

    Returns:
        bool: True if operation was successful, False otherwise
    """
    try:
        inspector = inspect(engine)

        # Check if table exists
        if inspector.has_table("research_log"):
            # For SQLite, we need to use raw SQL for DROP TABLE
            with engine.connect() as conn:
                conn.execute("DROP TABLE research_log")
                conn.commit()
            logger.info("Successfully removed redundant 'research_log' table")
            return True
        else:
            logger.info("Table 'research_log' does not exist, no action needed")
            return True
    except Exception:
        logger.exception("Error removing research_log table")
        return False


def create_research_strategy_table(engine):
    """
    Create the research_strategies table if it doesn't exist

    Args:
        engine: SQLAlchemy engine

    Returns:
        bool: True if operation was successful, False otherwise
    """
    try:
        inspector = inspect(engine)

        # Check if table exists
        if not inspector.has_table("research_strategies"):
            # Create the table using ORM
            Base.metadata.create_all(
                engine, tables=[ResearchStrategy.__table__]
            )
            logger.info("Successfully created 'research_strategies' table")
            return True
        else:
            logger.info(
                "Table 'research_strategies' already exists, no action needed"
            )
            return True
    except Exception:
        logger.exception("Error creating research_strategies table")
        return False


def create_benchmark_tables(engine):
    """
    Create benchmark tables if they don't exist

    Args:
        engine: SQLAlchemy engine

    Returns:
        bool: True if operation was successful, False otherwise
    """
    try:
        from .benchmark_schema import create_benchmark_tables_simple

        inspector = inspect(engine)

        # Check if benchmark tables already exist
        if not inspector.has_table("benchmark_runs"):
            # Create all benchmark tables using simple schema
            create_benchmark_tables_simple(engine)
            logger.info("Successfully created benchmark tables")
            return True
        else:
            logger.info("Benchmark tables already exist, no action needed")
            return True
    except Exception:
        logger.exception("Error creating benchmark tables")
        return False


def create_rate_limiting_tables(engine):
    """
    Create rate limiting tables if they don't exist

    Args:
        engine: SQLAlchemy engine

    Returns:
        bool: True if operation was successful, False otherwise
    """
    try:
        inspector = inspect(engine)

        tables_to_create = []

        # Check if rate_limit_attempts table exists
        if not inspector.has_table("rate_limit_attempts"):
            tables_to_create.append(RateLimitAttempt.__table__)
            logger.info("Need to create 'rate_limit_attempts' table")

        # Check if rate_limit_estimates table exists
        if not inspector.has_table("rate_limit_estimates"):
            tables_to_create.append(RateLimitEstimate.__table__)
            logger.info("Need to create 'rate_limit_estimates' table")

        if tables_to_create:
            # Create the tables using ORM
            Base.metadata.create_all(engine, tables=tables_to_create)
            logger.info(
                f"Successfully created {len(tables_to_create)} rate limiting tables"
            )
        else:
            logger.info("Rate limiting tables already exist, no action needed")

        return True
    except Exception:
        logger.exception("Error creating rate limiting tables")
        return False


def add_research_id_to_benchmark_results(engine):
    """
    Add research_id column to benchmark_results table if it doesn't exist.
    """
    try:
        import sqlite3

        # Get database path from engine
        db_path = engine.url.database

        logger.info("Checking if benchmark_results needs research_id column...")

        conn = sqlite3.connect(db_path)

        try:
            cursor = conn.cursor()

            # Check if table exists
            cursor.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='benchmark_results'"
            )
            if not cursor.fetchone():
                logger.info("benchmark_results table does not exist, skipping")
                return True

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

            if has_research_id:
                logger.info("benchmark_results already has research_id column")
                return True

            # Add research_id column
            logger.info("Adding research_id column to benchmark_results table")
            cursor.execute(
                "ALTER TABLE benchmark_results ADD COLUMN research_id TEXT"
            )

            conn.commit()
            logger.info(
                "Successfully added research_id column to benchmark_results"
            )
            return True

        finally:
            conn.close()

    except Exception:
        logger.exception("Error adding research_id column to benchmark_results")
        return False


def add_uuid_id_column_to_research_history(engine):
    """
    Adds a new `uuid_id` string column to the `research_history` table if it
    does not exist already.
    """
    try:
        import sqlite3

        # Get database path from engine
        db_path = engine.url.database

        logger.info("Checking if research_history needs uuid_id column...")

        conn = sqlite3.connect(db_path)

        try:
            cursor = conn.cursor()

            # Check if table exists
            cursor.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='research_history'"
            )
            if not cursor.fetchone():
                logger.info("research_history table does not exist, skipping")
                return True

            # Check if uuid_id column already exists
            cursor.execute("PRAGMA table_info(research_history)")
            columns = cursor.fetchall()
            has_uuid_id = any(col[1] == "uuid_id" for col in columns)

            if has_uuid_id:
                logger.info("research_history already has uuid_id column")
                return True

            # Add uuid_id column
            logger.info("Adding uuid_id column to research_history table")
            cursor.execute(
                "ALTER TABLE research_history ADD COLUMN uuid_id CHAR(36)"
            )

            conn.commit()
            logger.info("Successfully added uuid_id column to research_history")
            return True

        finally:
            conn.close()

    except Exception:
        logger.exception("Error adding uuid_id column to research_history")
        return False


def convert_research_id_to_string_if_needed(engine):
    """
    Convert research_id columns from Integer to String in all tables.
    Preserves existing data by converting integer IDs to string format.
    Only runs if integer research_id columns are detected.
    """
    try:
        import sqlite3

        # Get database path from engine
        db_path = engine.url.database

        logger.info(
            "Checking if research_id columns need conversion to string..."
        )

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

        try:
            cursor = conn.cursor()

            # List of tables that might have research_id columns
            tables_to_check = [
                "token_usage",
                "model_usage",
                "search_calls",
                "benchmark_results",
            ]

            tables_needing_conversion = []

            # Check which tables need conversion
            for table_name in tables_to_check:
                # Check if table exists
                cursor.execute(
                    "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
                    (table_name,),
                )
                if not cursor.fetchone():
                    continue

                # Check if research_id column exists and is integer type
                cursor.execute(f"PRAGMA table_info({table_name})")
                columns = cursor.fetchall()

                for col in columns:
                    col_name, col_type = col[1], col[2]
                    if col_name == "research_id" and (
                        "INTEGER" in col_type.upper()
                        or "INT" in col_type.upper()
                    ):
                        tables_needing_conversion.append(table_name)
                        break

            if not tables_needing_conversion:
                logger.info(
                    "All research_id columns are already string type, no conversion needed"
                )
                return True

            logger.info(
                f"Converting research_id to string in tables: {tables_needing_conversion}"
            )

            # Convert each table
            for table_name in tables_needing_conversion:
                logger.info(f"Converting {table_name} 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!"
            )
            return True

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

    except Exception:
        logger.exception("Error converting research_id columns to string")
        return False


def create_provider_models_table(engine):
    """Create provider_models table for caching available models"""
    with engine.connect() as conn:
        result = conn.execute(
            text(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='provider_models'"
            )
        )
        if result.fetchone():
            logger.info(
                "Table 'provider_models' already exists, no action needed"
            )
            return

        logger.info("Creating 'provider_models' table...")

        # Create the table
        conn.execute(
            text(
                """
            CREATE TABLE provider_models (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                provider VARCHAR(50) NOT NULL,
                model_key VARCHAR(255) NOT NULL,
                model_label VARCHAR(255) NOT NULL,
                model_metadata JSON,
                last_updated DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL,
                UNIQUE(provider, model_key)
            )
            """
            )
        )

        # Create index on provider
        conn.execute(
            text(
                "CREATE INDEX ix_provider_models_provider ON provider_models (provider)"
            )
        )

        conn.commit()
        logger.info("Table 'provider_models' created successfully")


def run_schema_upgrades():
    """
    Run all schema upgrade operations on the database

    Returns:
        bool: True if all upgrades successful, False otherwise
    """
    # Check if database exists
    if not os.path.exists(DB_PATH):
        logger.warning(
            f"Database not found at {DB_PATH}, skipping schema upgrades"
        )
        return False

    logger.info(f"Running schema upgrades on {DB_PATH}")

    try:
        # Create SQLAlchemy engine
        engine = create_engine(f"sqlite:///{DB_PATH}")

        # 1. Remove the redundant research_log table
        remove_research_log_table(engine)

        # 2. Create research_strategies table
        create_research_strategy_table(engine)

        # 3. Create benchmark tables
        create_benchmark_tables(engine)

        # 4. Create rate limiting tables
        create_rate_limiting_tables(engine)

        # 5. Add research_id column to benchmark_results if missing
        add_research_id_to_benchmark_results(engine)

        # 6. Convert research_id columns from integer to string
        convert_research_id_to_string_if_needed(engine)

        # 7. Add uuid_id column to research_history if missing
        add_uuid_id_column_to_research_history(engine)

        # 8. Create provider_models table for caching
        create_provider_models_table(engine)

        logger.info("Schema upgrades completed successfully")
        return True
    except Exception:
        logger.exception("Error during schema upgrades")
        return False


if __name__ == "__main__":
    run_schema_upgrades()
