"""
Migration script to add kb_type, priority, is_enabled, doc_count to knowledge_bases.
Run: python scripts/migrations/run_migration.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 knowledge_bases table."""
    settings = get_settings()
    engine = create_async_engine(settings.database_url, echo=True)
    async_session_maker = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

    migration_sql = """
    -- Add kb_type column
    ALTER TABLE knowledge_bases 
    ADD COLUMN IF NOT EXISTS kb_type VARCHAR DEFAULT 'general';

    -- Add priority column
    ALTER TABLE knowledge_bases 
    ADD COLUMN IF NOT EXISTS priority INTEGER DEFAULT 0;

    -- Add is_enabled column
    ALTER TABLE knowledge_bases 
    ADD COLUMN IF NOT EXISTS is_enabled BOOLEAN DEFAULT TRUE;

    -- Add doc_count column
    ALTER TABLE knowledge_bases 
    ADD COLUMN IF NOT EXISTS doc_count INTEGER DEFAULT 0;
    """

    async with async_session_maker() as session:
        for statement in migration_sql.strip().split(';'):
            statement = statement.strip()
            if statement and not statement.startswith('--'):
                try:
                    await session.execute(text(statement))
                    print(f"Executed: {statement[:50]}...")
                except Exception as e:
                    if "already exists" in str(e).lower() or "duplicate" in str(e).lower():
                        print(f"Skipped (already exists): {statement[:50]}...")
                    else:
                        raise
        
        await session.commit()
        print("\nMigration completed successfully!")

    await engine.dispose()


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