"""
Migration: Add provider column to TokenUsage table

This migration adds the provider column to the TokenUsage table to enable
accurate cost tracking based on both model and provider information.
"""

import logging
from pathlib import Path

from sqlalchemy import text

from .database import get_metrics_db

logger = logging.getLogger(__name__)


def add_provider_column_to_token_usage():
    """Add provider column to TokenUsage table."""
    try:
        db = get_metrics_db()

        with db.get_session() as session:
            # Check if provider column already exists
            result = session.execute(
                text(
                    """
                SELECT COUNT(*) as count
                FROM pragma_table_info('token_usage')
                WHERE name='provider'
            """
                )
            )

            provider_exists = result.fetchone()[0] > 0

            if provider_exists:
                logger.info(
                    "Provider column already exists in token_usage table"
                )
                return True

            logger.info("Adding provider column to token_usage table...")

            # Add the provider column
            session.execute(
                text(
                    """
                ALTER TABLE token_usage
                ADD COLUMN provider VARCHAR
            """
                )
            )

            # Try to populate provider info for existing records based on model name patterns
            logger.info("Populating provider info for existing records...")

            # Update known local model providers
            local_model_updates = [
                (
                    "ollama",
                    [
                        "ollama",
                        "llama",
                        "mistral",
                        "gemma",
                        "qwen",
                        "codellama",
                        "vicuna",
                        "alpaca",
                    ],
                ),
                ("openai", ["gpt-", "davinci", "curie", "babbage", "ada"]),
                ("anthropic", ["claude"]),
                ("google", ["gemini", "bard"]),
            ]

            for provider, model_patterns in local_model_updates:
                for pattern in model_patterns:
                    session.execute(
                        text(
                            """
                        UPDATE token_usage
                        SET provider = :provider
                        WHERE provider IS NULL
                        AND (LOWER(model_name) LIKE :pattern OR LOWER(model_name) LIKE :pattern_percent)
                    """
                        ),
                        {
                            "provider": provider,
                            "pattern": pattern,
                            "pattern_percent": f"%{pattern}%",
                        },
                    )

            # Set any remaining NULL providers to 'unknown'
            session.execute(
                text(
                    """
                UPDATE token_usage
                SET provider = 'unknown'
                WHERE provider IS NULL
            """
                )
            )

            session.commit()
            logger.info(
                "Successfully added provider column and populated existing data"
            )
            return True

    except Exception as e:
        logger.error(f"Error adding provider column to token_usage: {e}")
        return False


def run_migration():
    """Run the provider column migration."""
    logger.info("Starting migration: Add provider column to TokenUsage")

    success = add_provider_column_to_token_usage()

    if success:
        logger.info("Migration completed successfully")
    else:
        logger.error("Migration failed")

    return success


if __name__ == "__main__":
    # Allow running migration directly
    import sys

    # Add the project root to the path
    project_root = Path(__file__).parent.parent.parent.parent
    sys.path.insert(0, str(project_root))

    logging.basicConfig(level=logging.INFO)
    success = run_migration()

    if success:
        print("✅ Migration completed successfully")
        sys.exit(0)
    else:
        print("❌ Migration failed")
        sys.exit(1)
