"""Migration script to add Phase 1 enhanced token tracking fields."""

import sqlite3
from pathlib import Path

from loguru import logger

from ..utilities.db_utils import DB_PATH


def migrate_enhanced_tracking():
    """Add Phase 1 enhanced tracking columns to existing token_usage table."""

    if not Path(DB_PATH).exists():
        logger.info("Database doesn't exist yet, skipping migration")
        return

    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()

        # Check if token_usage table exists
        cursor.execute(
            """
            SELECT name FROM sqlite_master
            WHERE type='table' AND name='token_usage'
        """
        )

        if not cursor.fetchone():
            logger.info(
                "token_usage table doesn't exist yet, skipping migration"
            )
            conn.close()
            return

        # Check if enhanced columns already exist
        cursor.execute("PRAGMA table_info(token_usage)")
        columns = [column[1] for column in cursor.fetchall()]

        # Define new columns to add
        new_columns = [
            ("research_query", "TEXT"),
            ("research_mode", "TEXT"),
            ("research_phase", "TEXT"),
            ("search_iteration", "INTEGER"),
            ("response_time_ms", "INTEGER"),
            ("success_status", "TEXT DEFAULT 'success'"),
            ("error_type", "TEXT"),
            ("search_engines_planned", "TEXT"),
            ("search_engine_selected", "TEXT"),
        ]

        # Add missing columns
        for column_name, column_type in new_columns:
            if column_name not in columns:
                logger.info(f"Adding column {column_name} to token_usage table")
                cursor.execute(
                    f"ALTER TABLE token_usage ADD COLUMN {column_name} {column_type}"
                )

        conn.commit()
        conn.close()

        logger.info("Enhanced token tracking migration completed successfully")

    except Exception as e:
        logger.exception(f"Error during enhanced token tracking migration: {e}")
        if "conn" in locals():
            conn.close()
        raise


if __name__ == "__main__":
    migrate_enhanced_tracking()
