#!/usr/bin/env python3
"""
Production Migration Script for llm -> llm_catalog_rel.
This script handles:
1. Creating the new relationship table if it doesn't exist.
2. Migrating existing single-type associations to the new table.
3. Dropping the old column (optional but recommended).

Usage:
    cd /home/hermesai/repos/sage
    ./py3/bin/python migrate_llm_catelog_prod.py
"""
import asyncio
import sys
sys.path.insert(0, '.')

from sqlor.dbpools import DBPools
from appPublic.jsonConfig import getConfig
from appPublic.uniqueID import getID

async def migrate():
    config = getConfig('.')
    db = DBPools(config.databases)
    
    # Sage 数据库
    dbname = 'sage'
    
    async with db.sqlorContext(dbname) as sor:
        # 1. 检查列是否存在
        check_col_sql = "select count(*) as cnt from information_schema.columns where table_schema = database() and table_name = 'llm' and column_name = 'llmcatelogid'"
        res = await sor.sqlExe(check_col_sql, {})
        has_column = res[0].cnt > 0 if res else False

        if not has_column:
            print("[SKIP] Column 'llmcatelogid' does not exist in 'llm' table. Migration not needed.")
            return True

        # 2. 创建表 (如果不存在)
        print("[STEP 1] Ensuring table 'llm_catalog_rel' exists...")
        create_sql = """
        CREATE TABLE IF NOT EXISTS llm_catalog_rel (
            id VARCHAR(32) NOT NULL PRIMARY KEY,
            llmid VARCHAR(32) NOT NULL,
            llmcatelogid VARCHAR(32) NOT NULL,
            UNIQUE KEY uq_llm_catelog (llmid, llmcatelogid),
            INDEX idx_llm (llmid),
            INDEX idx_catelog (llmcatelogid)
        )
        """
        try:
            await sor.sqlExe(create_sql, {})
            print("  -> Table created or exists.")
        except Exception as e:
            print(f"  -> Table creation warning: {e}")

        # 2.1 确保唯一索引存在
        print("[STEP 1.1] Ensuring unique index 'uq_llm_catelog' exists...")
        try:
            await sor.sqlExe("ALTER TABLE llm_catalog_rel ADD UNIQUE INDEX uq_llm_catelog (llmid, llmcatelogid)", {})
            print("  -> Unique index added.")
        except Exception as e:
            if "Duplicate key name" in str(e) or "Multiple primary key defined" in str(e):
                print("  -> Unique index already exists.")
            else:
                print(f"  -> Index creation warning (may already exist): {e}")

        # 3. 迁移数据
        print("[STEP 2] Migrating data...")
        sql = "select id, llmcatelogid from llm where llmcatelogid is not null and llmcatelogid != ''"
        rows = await sor.sqlExe(sql, {})
        print(f"  -> Found {len(rows)} records to migrate.")
        
        migrated_count = 0
        error_count = 0
        
        for r in rows:
            # Check if already migrated to avoid duplicates if re-running
            check_rel_sql = "select count(*) as cnt from llm_catalog_rel where llmid = ${llmid}$ and llmcatelogid = ${catelogid}$"
            rel_res = await sor.sqlExe(check_rel_sql, {'llmid': r['id'], 'catelogid': r['llmcatelogid']})
            if rel_res and rel_res[0].cnt > 0:
                continue

            new_id = getID()
            data = {
                'id': new_id,
                'llmid': r['id'],
                'llmcatelogid': r['llmcatelogid']
            }
            try:
                await sor.C('llm_catalog_rel', data)
                migrated_count += 1
            except Exception as e:
                print(f"  -> Error migrating {r['id']}: {e}")
                error_count += 1

        print(f"  -> Migration complete. Inserted: {migrated_count}, Errors: {error_count}")

        # 4. 删除旧列 (生产环境谨慎执行，建议确认迁移成功后手动执行)
        # print("[STEP 3] Dropping old column 'llmcatelogid'...")
        # try:
        #     await sor.sqlExe("alter table llm drop column llmcatelogid", {})
        #     print("  -> Column dropped.")
        # except Exception as e:
        #     print(f"  -> Drop column error: {e}")

        print("\n[MIGRATION DONE]")
        print("Please verify data in 'llm_catalog_rel'.")
        print("If correct, you may manually run: alter table llm drop column llmcatelogid;")

    return True

if __name__ == '__main__':
    print("Starting production migration...")
    try:
        success = asyncio.get_event_loop().run_until_complete(migrate())
        if success:
            sys.exit(0)
        else:
            sys.exit(1)
    except Exception as e:
        print(f"FATAL ERROR: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)