"""
Migration script to add monitoring fields to chat_messages table.
Run: python scripts/migrations/add_chat_message_fields.py
"""

import asyncio
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlmodel.ext.asyncio.session import AsyncSession

from app.core.config import get_settings


async def run_migration():
    """Run the migration to add new columns to chat_messages table."""
    settings = get_settings()
    engine = create_async_engine(settings.database_url, echo=True)
    async_session_maker = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

    statements = [
        "ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS prompt_template_id UUID",
        "ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS intent_rule_id UUID",
        "ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS flow_instance_id UUID",
        "ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS guardrail_triggered BOOLEAN DEFAULT FALSE",
        "ALTER TABLE chat_messages ADD COLUMN IF NOT EXISTS guardrail_words JSONB",
        "CREATE INDEX IF NOT EXISTS ix_chat_messages_tenant_template ON chat_messages(tenant_id, prompt_template_id)",
        "CREATE INDEX IF NOT EXISTS ix_chat_messages_tenant_intent ON chat_messages(tenant_id, intent_rule_id)",
        "CREATE INDEX IF NOT EXISTS ix_chat_messages_tenant_flow ON chat_messages(tenant_id, flow_instance_id)",
    ]

    async with async_session_maker() as session:
        for statement in statements:
            try:
                await session.execute(text(statement))
                print(f"Executed: {statement[:60]}...")
            except Exception as e:
                error_str = str(e).lower()
                if "already exists" in error_str or "duplicate" in error_str:
                    print(f"Skipped (already exists): {statement[:60]}...")
                else:
                    print(f"Error: {e}")
                    raise
        
        await session.commit()
        print("\nMigration completed successfully!")

    await engine.dispose()


if __name__ == "__main__":
    asyncio.run(run_migration())
