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

# 1. 在 import 区域加 AsyncLocalStorage
if 'AsyncLocalStorage' not in src:
    # 找一个 import 区块
    src = src.replace(
        "import { AsyncLocalStorage } from 'async_hooks';",
        "import { AsyncLocalStorage } from 'async_hooks';",
        0
    )
    # 实际上还没 import，需要在文件头部加
    # 找到第一个 import 行
    import_line = "import { AsyncLocalStorage } from 'async_hooks';\n"
    if 'import { AsyncLocalStorage' not in src:
        # 找第一个 import 行的位置之前插入
        first_import = src.find('import ')
        src = src[:first_import] + import_line + src[first_import:]

# 2. 在 findVendorForModel 之后加 ALS helper + resolveNodeForModel + pickNextModelAfterFailure
# 找位置：findVendorForModel 函数体结束
helper_code = """
// ============ HA: AsyncLocalStorage 传递 tried nodes ============
const _triedNodesStorage = new AsyncLocalStorage<Set<string>>();
function _runWithTriedNodes<T>(triedSet: Set<string>, fn: () => T): T {
  return _triedNodesStorage.run(triedSet, fn);
}
function _currentTriedNodes(): Set<string> {
  return _triedNodesStorage.getStore() || new Set();
}

/** 根据 modelId + tried set 找到下一个未尝试的 provider node */
function resolveNodeForModel(modelId: string, tried: Map<string, any>): any | null {
  const registry = 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;
}

/** 从当前 node 找下一个未 tried 的 node, 返回 next modelId */
function pickNextModelAfterFailure(currentModelId: string, tried: Map<string, any>, currentNode: any): string | null {
  const registry = getLlmRegistry();
  const available = registry.listAvailable();
  if (available.length === 0) return null;
  const startIndex = currentNode ? (available.indexOf(currentNode) + 1) % available.length : 0;
  for (let i = 0; i < available.length; i++) {
    const idx = (startIndex + i) % available.length;
    const node = available[idx];
    if (node === currentNode) continue;
    const fp = node.provider.vendor + ':' + (n.provider._apiKey || '').slice(-12);
    if (tried.has(fp)) continue;
    if (node.provider.hasModel(currentModelId)) {
      console.log(`[LLM] 同模型切换到供应商 ${node.provider.displayName}，模型 ${currentModelId}`);
      return currentModelId;
    }
    for (const modelId of node.provider.textModels) {
      const cfg = node.provider.getModelConfig(modelId);
      if (cfg?.canonicalModel === currentModelId) {
        console.log(`[LLM] 别名模型切换到供应商 ${node.provider.displayName}，模型 ${modelId}`);
        return modelId;
      }
    }
    if (node.provider.textModels.length > 0) {
      const nextId = node.provider.textModels[0];
      console.log(`[LLM] 切换到供应商 ${node.provider.displayName}，模型 ${nextId}`);
      return nextId;
    }
  }
  return null;
}
"""

# fix typo: (n.provider._apiKey should be (node.provider._apiKey
helper_code = helper_code.replace("(n.provider._apiKey || '')", "(node.provider._apiKey || '')")

# 插在 findVendorForModel 之后
if 'resolveNodeForModel' not in src:
    # 找 findVendorForModel 结束的位置
    fvf_marker = "function findVendorForModel"
    idx = src.find(fvf_marker)
    if idx < 0:
        print("ERROR: findVendorForModel not found")
        raise SystemExit(1)
    # 找函数体结束的 '}'（找匹配的最外层）
    depth = 0
    i = src.find('{', idx)
    end = -1
    while i < len(src):
        if src[i] == '{': depth += 1
        elif src[i] == '}':
            depth -= 1
            if depth == 0:
                end = i + 1
                break
        i += 1
    if end < 0:
        print("ERROR: findVendorForModel body not found")
        raise SystemExit(1)
    src = src[:end] + "\n" + helper_code + src[end:]
    print("OK: helpers injected after findVendorForModel")

# 3. 改 findVendorForModel 跳过 tried
OLD_FVF = """function findVendorForModel(modelId) {
    const registry = getLlmRegistry();
    return registry.listAvailable().find(n => n.provider.hasModel(modelId));
}"""

NEW_FVF = """function findVendorForModel(modelId) {
    const registry = getLlmRegistry();
    const tried = _currentTriedNodes();
    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;
    }
    return registry.listAvailable().find(n => n.provider.hasModel(modelId)) || null;
}"""

if OLD_FVF in src:
    src = src.replace(OLD_FVF, NEW_FVF)
    print("OK: findVendorForModel patched")
else:
    print("WARN: findVendorForModel old pattern not found")

# 4. 改 invokeWithRetry (v3 整段)
# 暂不改 src invokeWithRetry — dist 的 v3 在跑, src 等下次有精力再合
# 但要保证 tsc 不报错
# 实际 tsc 不会编译 dist 用的代码到 dist, 所以 dist 的 v3 是 node 跑时直接要的
# 跑 npm run build 才会重编 src → 覆盖 dist
# 所以我们必须改 src invokeWithRetry

# 找 invokeWithRetry 在 src 的位置
inv_marker = "async function invokeWithRetry"
idx = src.find(inv_marker)
print(f"invokeWithRetry starts at line {src[:idx].count(chr(10))+1}")

# 看 src 当前 invokeWithRetry 大致行数
# 直接 find 函数结束
depth = 0
i = src.find('{', idx)
end = -1
while i < len(src):
    if src[i] == '{': depth += 1
    elif src[i] == '}':
        depth -= 1
        if depth == 0:
            end = i + 1
            break
    i += 1
if end < 0:
    print("ERROR: invokeWithRetry body not found")
    raise SystemExit(1)

# print(src[idx:end][:500])
# print("---")
# print(src[idx:end][-200:])

# 替换为 v3 版本
V3 = """async function invokeWithRetry(
  fn: (modelId: string) => Promise<any>,
  modelId: string,
  callType: string = 'llm_chat',
): Promise<any> {
  const tried = new Map<string, any>();
  let cur = modelId;
  let lastErr: any = null;

  for (let round = 0; round < HA_RETRY_CONFIG.maxFallbackRounds; round++) {
    // 关键: 选下一个未 tried 的 node (而不是用 resolveProviderKey 取第一个匹配)
    const node = resolveNodeForModel(cur, tried);
    if (!node) {
      console.warn(`[LLM-HA] 已尝试所有 node，停止 fallback`);
      break;
    }
    const fp = node.provider.vendor + ':' + (node.provider._apiKey || '').slice(-12);
    tried.set(fp, { model: cur });
    let providerError: any = null;

    for (let attempt = 0; attempt < HA_RETRY_CONFIG.maxRetriesPerProvider; attempt++) {
      try {
        const result = await _runWithTriedNodes(new Set(tried.keys()), () => fn(cur));
        return result;
      } catch (error: any) {
        const em = error?.message || String(error);
        lastErr = error;
        providerError = error;
        if (isBusinessError(em) && !isNetworkError(em)) {
          console.error(`[LLM-HA] ${fp}/${cur} 业务错误，立即抛出: ${em.substring(0, 100)}`);
          throw error;
        }
        if (isNetworkError(em)) {
          console.warn(`[LLM-HA] ${fp}/${cur} 网络错误，立即切换: ${em.substring(0, 100)}`);
          break;
        }
        if (attempt < HA_RETRY_CONFIG.maxRetriesPerProvider - 1) {
          const delay = backoffWithJitter(attempt);
          console.log(`[LLM-HA] ${fp}/${cur} 第${attempt + 1}次失败: ${em.substring(0, 100)}, ${delay}ms后重试`);
          await new Promise(r => setTimeout(r, delay));
        }
      }
    }

    const em = providerError?.message || 'unknown';
    const et = isNetworkError(em) ? 'network' : (isBusinessError(em) ? 'business' : 'unknown');
    tried.set(fp, { model: cur, error: em, attempts: HA_RETRY_CONFIG.maxRetriesPerProvider, errorType: et });
    console.log(`[LLM-HA] ${fp}/${cur} 失败 (${et})，尝试切换供应商`);

    const next = pickNextModelAfterFailure(cur, tried, node);
    if (!next) {
      console.warn(`[LLM-HA] 找不到下一个未尝试的供应商`);
      break;
    }
    cur = next;
    console.log(`[LLM-HA] 切换到 ${cur}`);
  }

  throw new AllProvidersFailedError(
    Array.from(tried.entries()).map(([provider, info]) => ({ provider, ...info }))
  );
}"""

# 找到调用 trySwitchModel 的位置，确认范围
# 用更可靠的方法：找 "async function invokeWithRetry" 到 "throw new AllProvidersFailedError" 的最外层
import re
m = re.search(r'async function invokeWithRetry[^{]*\{', src)
if not m:
    print("ERROR: invokeWithRetry def not found")
    raise SystemExit(1)
start = m.start()
# 用括号配对找结束
depth = 0
i = src.find('{', start)
end = -1
while i < len(src):
    if src[i] == '{': depth += 1
    elif src[i] == '}':
        depth -= 1
        if depth == 0:
            end = i + 1
            break
    i += 1
if end < 0:
    print("ERROR: invokeWithRetry body end not found")
    raise SystemExit(1)

# 也包含 HA_RETRY_CONFIG / AllProvidersFailedError 等需要在文件前部定义
# 检查它们是否已存在
if 'HA_RETRY_CONFIG' not in src:
    print("ERROR: HA_RETRY_CONFIG not defined in src - run previous patch first")
    raise SystemExit(1)
if 'class AllProvidersFailedError' not in src:
    print("ERROR: AllProvidersFailedError not defined in src - run previous patch first")
    raise SystemExit(1)

old_func = src[start:end]
new_func = V3
src = src[:start] + new_func + src[end:]
print(f"OK: invokeWithRetry replaced ({len(old_func)} -> {len(new_func)} chars)")

# 5. 把 HA_RETRY_CONFIG 改为 export const (dist 也是 exports.HA_RETRY_CONFIG, src 可能要 export)
# 实际 tsc 编译 src -> dist 时, exports.X 变成 module.exports.X, 正常
# 但 HA_RETRY_CONFIG 在 src 是 const 吗? 让我检查
# 既然前面 tsc 没报错, 应该是 const 形式

with open(PATH, 'w', encoding='utf-8') as f:
    f.write(src)
print("OK: src/services/llm/index.ts fully synced with v3")