#!/usr/bin/env python3
"""pipeline-llm 初始化迁移：旧 `llm` 表 → 模型治理新表（幂等，可重跑）。

背景（2026-09-04）：产线平台模型调用统一收敛到模型治理模块推理 API，
旧 llm 表停用。本脚本把存量模型迁到新表并给每个机构配好容错策略，
切换后立即可用（「初始化设置一个可用的模型」）。

迁移规则：
  1. 旧行按 (provider, api_base) 归并为供应商（llm_vendor，端点目录一条）
  2. 同供应商下按 api_key 明文归并为账号（llm_account，AES 重加密）
     —— RC4(password_key) 解密旧 key → AES(password_key) 加密入新表
  3. 每旧行一个模型（llm_model，name 保持原名，存量引用不漂移；
     capabilities 'text'→'t2t'，其余照能力类型标准）
  4. 给 '0' + users 表全部机构配组织容错策略：主模型 = 默认 t2t
     （优先 qwen3.8-max，否则首个 t2t），备链 = 其余 t2t 模型
  5. 定价初始 0（未定价模型调用不冻结额度；接入定价后再配）

用法（应用根目录）：
    py3/bin/python pkgs/pipeline-llm/scripts/migrate_legacy_llm.py [--dry-run]
"""
import asyncio
import json
import os
import sys

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
# 应用根：本脚本位于 <app>/pkgs/pipeline-llm/scripts/，上溯三级
APP_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(SCRIPT_DIR)))
sys.path.insert(0, os.path.join(APP_ROOT, 'py3', 'lib', 'python3.10', 'site-packages'))
sys.path.insert(0, APP_ROOT)

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

CAP_MAP = {'text': 't2t', '': 't2t'}   # 旧 capabilities → 能力类型标准


def _decrypt_legacy(enc, password_key):
    """旧表 api_key 是 RC4(password_key) 加密。解不出原样返回。"""
    if not enc:
        return ''
    try:
        from appPublic.rc4 import unpassword
        return unpassword(enc, password_key)
    except Exception:
        return enc


def _encrypt_new(plain, password_key):
    """新表 api_key 用 AES 单层加密（与模块 gateway.encrypt_api_key 一致）。"""
    if not plain:
        return ''
    from appPublic.aes import aes_encode_b64
    return aes_encode_b64(password_key, plain)


async def _get_or_create_vendor(sor, provider, base_url, dry):
    name = (provider or '').strip() or '未知供应商'
    recs = await sor.sqlExe(
        "SELECT id, endpoints FROM llm_vendor WHERE name=${n}$ AND status='active' LIMIT 1",
        {"n": name})
    await sor.sqlExe("COMMIT", {})
    if recs:
        vid = getattr(recs[0], 'id', '')
        # 端点目录缺此 base_url 则补
        eps = []
        try:
            eps = json.loads(getattr(recs[0], 'endpoints', '') or '[]') or []
        except Exception:
            eps = []
        if not any((e.get('base_url') or '') == base_url for e in eps):
            eps.append({'base_url': base_url, 'region': 'domestic', 'timeout': 60})
            if not dry:
                await sor.sqlExe(
                    "UPDATE llm_vendor SET endpoints=${e}$ WHERE id=${i}$",
                    {"e": json.dumps(eps, ensure_ascii=False), "i": vid})
                await sor.sqlExe("COMMIT", {})
        idx = next(i for i, e in enumerate(eps) if (e.get('base_url') or '') == base_url)
        return vid, idx
    vid = getID()
    eps = [{'base_url': base_url, 'region': 'domestic', 'timeout': 60}]
    if not dry:
        await sor.C('llm_vendor', {
            'id': vid, 'name': name,
            'endpoints': json.dumps(eps, ensure_ascii=False),
            'description': '由旧 llm 表迁移生成', 'status': 'active', 'org_id': '0',
        })
        await sor.sqlExe("COMMIT", {})
    return vid, 0


async def _get_or_create_account(sor, vendor_id, ep_idx, key_plain, password_key, dry):
    """同供应商同 key 复用账号（余额/用量延续）。按解密后明文比对。"""
    recs = await sor.sqlExe(
        "SELECT id, api_key FROM llm_account WHERE vendor_id=${v}$", {"v": vendor_id})
    await sor.sqlExe("COMMIT", {})
    for r in (recs or []):
        try:
            from appPublic.aes import aes_decode_b64
            if aes_decode_b64(password_key, getattr(r, 'api_key', '') or '') == key_plain:
                return getattr(r, 'id', '')
        except Exception:
            continue
    aid = getID()
    if not dry:
        await sor.C('llm_account', {
            'id': aid, 'vendor_id': vendor_id,
            'name': '迁移账号-%s' % (key_plain[:6] if key_plain else 'nokey'),
            'api_key': _encrypt_new(key_plain, password_key),
            'endpoint_ids': json.dumps([ep_idx]),
            'balance': 0, 'total_recharge': 0, 'status': 'active', 'org_id': '0',
        })
        await sor.sqlExe("COMMIT", {})
    return aid


async def main(dry=False):
    config = getConfig(APP_ROOT, NS={'workdir': APP_ROOT, 'ProgramPath': ProgramPath()})
    DBPools(config.databases)
    password_key = config.password_key or 'QRIVSRHrthhwyjy176556332'

    async with DBPools().sqlorContext('pipeline') as sor:
        recs = await sor.sqlExe(
            "SELECT id, name, provider, model_id, api_base, api_key, capabilities, org_id "
            "FROM llm WHERE status='active' ORDER BY name", {})
        await sor.sqlExe("COMMIT", {})
        if not recs:
            print("旧 llm 表无 active 模型，无需迁移")
            return

        print("旧表 active 模型 %d 个，开始迁移（%s）..." % (len(recs), 'dry-run' if dry else '实写'))
        created = {'vendor': 0, 'account': 0, 'model': 0, 'policy': 0}
        t2t_model_ids = []      # (id, name)
        for r in recs:
            name = getattr(r, 'name', '') or ''
            api_base = (getattr(r, 'api_base', '') or '').rstrip('/')
            cap = CAP_MAP.get((getattr(r, 'capabilities', '') or '').strip().lower(),
                              (getattr(r, 'capabilities', '') or '').strip().lower() or 't2t')
            key_plain = _decrypt_legacy(getattr(r, 'api_key', '') or '', password_key)
            if key_plain and key_plain == getattr(r, 'api_key', ''):
                print("  ⚠ %s: api_key 解密结果与密文相同（可能已是明文或解密失败），原样迁移" % name)

            vid, ep_idx = await _get_or_create_vendor(sor, getattr(r, 'provider', ''), api_base, dry)
            await _get_or_create_account(sor, vid, ep_idx, key_plain, password_key, dry)

            exist = await sor.sqlExe("SELECT id FROM llm_model WHERE name=${n}$", {"n": name})
            await sor.sqlExe("COMMIT", {})
            if exist:
                mid = getattr(exist[0], 'id', '')
                print("  = 模型已存在跳过: %s" % name)
            else:
                mid = getID()
                if not dry:
                    await sor.C('llm_model', {
                        'id': mid, 'vendor_id': vid,
                        'name': name,
                        'vendor_model_id': getattr(r, 'model_id', '') or name,
                        'capability': cap, 'sync_mode': 'sync',
                        'profile_id': '', 'query_profile_ids': '',
                        'ppid': '',
                        'default_params': '', 'status': 'active',
                        'description': '由旧 llm 表迁移生成', 'org_id': '0',
                    })
                    await sor.sqlExe("COMMIT", {})
                created['model'] += 1
                print("  + 模型: %s (cap=%s vendor=%s)" % (name, cap, getattr(r, 'provider', '')))
            if cap == 't2t':
                t2t_model_ids.append((mid, name))

        if not t2t_model_ids:
            print("⚠ 无 t2t 模型，跳过策略配置")
            return

        # 默认主模型：优先 max 档（能力最强），否则名字序第一个
        primary = None
        for mid, name in t2t_model_ids:
            if 'max' in name.lower():
                primary = (mid, name)
                break
        primary = primary or t2t_model_ids[0]
        backups = [mid for mid, name in t2t_model_ids if mid != primary[0]]
        print("策略：主模型=%s，备链=%d 个" % (primary[1], len(backups)))

        # 给系统级 + 全部机构配策略（策略即开关：配了才走治理）
        orgs = ['0']
        urecs = await sor.sqlExe("SELECT DISTINCT orgid FROM users WHERE orgid IS NOT NULL", {})
        await sor.sqlExe("COMMIT", {})
        for u in (urecs or []):
            o = getattr(u, 'orgid', '') or ''
            if o and o not in orgs:
                orgs.append(o)
        for org in orgs:
            prec = await sor.sqlExe(
                "SELECT id FROM llm_org_policy WHERE org_id=${o}$", {"o": org})
            await sor.sqlExe("COMMIT", {})
            if prec:
                print("  = 策略已存在跳过: org=%s" % org)
                continue
            if not dry:
                await sor.C('llm_org_policy', {
                    'id': getID(), 'org_id': org,
                    'primary_model_id': primary[0],
                    'backup_model_ids': json.dumps(backups, ensure_ascii=False),
                    'endpoint_pref': 'any', 'status': 'active',
                })
                await sor.sqlExe("COMMIT", {})
            created['policy'] += 1
            print("  + 策略: org=%s" % org)

        print("\n迁移完成：%s" % json.dumps(created, ensure_ascii=False))
        print("注意：governance_enabled 有进程内缓存，部署迁移后须重启服务生效。")


if __name__ == '__main__':
    dry = '--dry-run' in sys.argv
    asyncio.run(main(dry))