import asyncio
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. Create table if not exists
        print("Creating llm_catalog_rel table...")
        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,
            INDEX idx_llm (llmid),
            INDEX idx_catelog (llmcatelogid)
        )
        """
        try:
            await sor.sqlExe(create_sql, {})
            print("Table llm_catalog_rel created or exists.")
        except Exception as e:
            print(f"Create table warning: {e}")

        # 2. Migrate data
        print("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 = 0
        for r in rows:
            new_id = getID()
            data = {
                'id': new_id,
                'llmid': r['id'],
                'llmcatelogid': r['llmcatelogid']
            }
            try:
                await sor.C('llm_catalog_rel', data)
                migrated += 1
            except Exception as e:
                print(f"Insert error for {r['id']}: {e}")

        print(f"Migration complete. Migrated {migrated} records.")

        # 3. Drop column
        print("Dropping column llmcatelogid from llm...")
        try:
            await sor.sqlExe("alter table llm drop column llmcatelogid", {})
            print("Column dropped.")
        except Exception as e:
            print(f"Drop column error: {e}")

    return True

if __name__ == '__main__':
    success = asyncio.get_event_loop().run_until_complete(migrate())
    if success:
        print('Migration complete.')
    else:
        print('Migration failed.')