PATH = '/data/ai/audio/server/dist/services/llm/index.js'
with open(PATH, 'r', encoding='utf-8') as f:
    src = f.read()

# 1. 注入 AsyncLocalStorage require (在 require 区块)
need_als = "const { AsyncLocalStorage } = require('async_hooks');" in src or 'AsyncLocalStorage' in src
if not need_als:
    # 找到 require('async_hooks') 或在顶部 require 后插入
    if "require('async_hooks')" not in src:
        # 在文件最顶部 require 区域插入
        src = src.replace(
            "Object.defineProperty(exports, '__esModule', { value: true });",
            "Object.defineProperty(exports, '__esModule', { value: true });\nconst { AsyncLocalStorage: _ALS } = require('async_hooks');\nconst _triedNodesStorage = new _ALS();\nfunction _runWithTriedNodes(triedSet, fn) { return _triedNodesStorage.run(triedSet, fn); }\nfunction _currentTriedNodes() { return _triedNodesStorage.getStore() || new Set(); }",
            1
        )

# 2. 改造 findVendorForModel：跳过 tried 节点
OLD_FIND = """function findVendorForModel(modelId) {
    const registry = (0, provider_registry_1.getLlmRegistry)();
    return registry.listAvailable().find(n => n.provider.hasModel(modelId));
}"""

NEW_FIND = """function findVendorForModel(modelId) {
    const registry = (0, provider_registry_1.getLlmRegistry)();
    const tried = _currentTriedNodes();
    // 关键: 跳过当前 call scope 里已 tried 的 node (避免 minimax Key1/Key2 同 modelId 时永远命中第一个)
    for (const n of registry.listAvailable()) {
        if (!n.provider.hasModel(modelId)) continue;
        const fp = n.provider.vendor + ':' + (n.provider._apiKey || '').slice(-12);
        if (tried.has(fp)) continue;
        return n;
    }
    // 全部 tried 过, 返回第一个匹配的 (会触发上层 break)
    return registry.listAvailable().find(n => n.provider.hasModel(modelId)) || null;
}"""

if OLD_FIND not in src:
    print("ERROR: findVendorForModel pattern not found")
    raise SystemExit(1)
src = src.replace(OLD_FIND, NEW_FIND)

# 3. 改造 invokeWithRetry：用 AsyncLocalStorage 包裹 fn 调用
# 找到 "for (let attempt = 0; attempt < exports.HA_RETRY_CONFIG.maxRetriesPerProvider; attempt++)" 这一段
OLD_RETRY_LOOP = """        for (let attempt = 0; attempt < exports.HA_RETRY_CONFIG.maxRetriesPerProvider; attempt++) {
            try {
                return await fn(cur);
            }"""

NEW_RETRY_LOOP = """        for (let attempt = 0; attempt < exports.HA_RETRY_CONFIG.maxRetriesPerProvider; attempt++) {
            // 用 AsyncLocalStorage 传递 tried set, 让 findVendorForModel 跳过已试 node
            const node = resolveNodeForModel(cur, tried);
            if (!node) break;
            const fp = node.provider.vendor + ':' + (node.provider._apiKey || '').slice(-12);
            tried.set(fp, { model: cur });
            try {
                const result = await _runWithTriedNodes(new Set(tried.keys()), () => fn(cur));
                return result;
            }"""

if OLD_RETRY_LOOP not in src:
    print("ERROR: retry loop pattern not found")
    raise SystemExit(1)
src = src.replace(OLD_RETRY_LOOP, NEW_RETRY_LOOP)

# 4. 加 resolveNodeForModel 函数（如果不存在）
if 'function resolveNodeForModel' not in src:
    HELPER = """
/** 根据 modelId + tried set 找到下一个未尝试的 provider node */
function resolveNodeForModel(modelId, tried) {
    const registry = (0, provider_registry_1.getLlmRegistry)();
    for (const n of registry.listAvailable()) {
        if (!n.provider.hasModel(modelId)) continue;
        const fp = n.provider.vendor + ':' + (n.provider._apiKey || '').slice(-12);
        if (!tried.has(fp)) return n;
    }
    return null;
}
"""
    # 插在 resolveProviderKey 之后
    src = src.replace("function resolveProviderKey(modelId) {", HELPER + "\nfunction resolveProviderKey(modelId) {", 1)

# 5. 简化 invokeWithRetry 的循环头（移除老的 tried.has(provider) 检查，因为现在用 resolveNodeForModel 决定）
# 找到 "if (tried.has(provider)) break;"
OLD_BREAK = """        const provider = resolveProviderKey(cur);
        if (tried.has(provider))
            break;"""

NEW_BREAK = """        const provider = resolveProviderKey(cur);
        // tried check is now done in resolveNodeForModel below; keep this as a guard
        if (tried.has(provider))
            continue;"""

if OLD_BREAK not in src:
    print("WARN: tried.has(provider) break pattern not found, skipping")
else:
    src = src.replace(OLD_BREAK, NEW_BREAK)

with open(PATH, 'w', encoding='utf-8') as f:
    f.write(src)
print("OK: full HA chain fix applied")