#!/usr/bin/env python3
"""
Database migration script to add call stack tracking columns to token_usage table.
This adds the Phase 1 call stack tracking functionality.
"""

import sqlite3
import sys
from pathlib import Path

from loguru import logger


def migrate_call_stack_tracking(db_path: str):
    """Add call stack tracking columns to the token_usage table.

    Args:
        db_path: Path to the SQLite database file
    """
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

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

        # Add call stack tracking columns if they don't exist
        new_columns = [
            ("calling_file", "TEXT"),
            ("calling_function", "TEXT"),
            ("call_stack", "TEXT"),
        ]

        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}"
                )
            else:
                logger.info(
                    f"Column {column_name} already exists in token_usage table"
                )

        conn.commit()
        logger.success(
            "Call stack tracking columns migration completed successfully"
        )

    except sqlite3.Error as e:
        logger.error(f"Database error during call stack migration: {e}")
        raise
    except Exception as e:
        logger.error(f"Unexpected error during call stack migration: {e}")
        raise
    finally:
        if conn:
            conn.close()


def find_database_file():
    """Find the metrics database file."""
    # Common locations for the database
    possible_paths = [
        "data/metrics.db",
        "../data/metrics.db",
        "../../data/metrics.db",
    ]

    for path in possible_paths:
        db_path = Path(path)
        if db_path.exists():
            return str(db_path.absolute())

    return None


if __name__ == "__main__":
    logger.info("Starting call stack tracking migration...")

    # Check if database path provided as argument
    if len(sys.argv) > 1:
        db_path = sys.argv[1]
    else:
        db_path = find_database_file()

    if not db_path:
        logger.error("Could not find metrics database file.")
        logger.info("Please provide the database path as an argument:")
        logger.info("python migrate_call_stack_tracking.py /path/to/metrics.db")
        sys.exit(1)

    if not Path(db_path).exists():
        logger.error(f"Database file does not exist: {db_path}")
        sys.exit(1)

    logger.info(f"Using database: {db_path}")

    try:
        migrate_call_stack_tracking(db_path)
        logger.success("Call stack tracking migration completed!")
    except Exception as e:
        logger.error(f"Migration failed: {e}")
        sys.exit(1)
