"""LLM 传输健壮性层(从 core/loop.py 析出,2026-07-23)。

关注点:provider wire 层的瞬态故障检测与自愈 —— 与 agent 控制流(ReAct 循环 /
工具执行 / 熔断)正交。收在这里的东西回答同一个问题:「这一轮 LLM 响应能不能用,
不能用怎么救」:

- 检测:畸形 arguments(JSON 解析失败)/ 必填 key 被吞(解析成功但键被流式乱序
  吞掉)/ 空响应(tc 空且正文空)/ finish_reason
- 留痕:三类故障各自 stdout + usage_events 双写(留痕绝不打断重试主路径)
- 重试策略 robust_stream:首败即降级非流式(同轮流式失败强相关,provider 服务端
  拼 tool_calls 绕开 delta 错位),salvage 可救则当轮继续
- usage/delta 提取:provider 差异归一

依赖注入纪律:取流的两条路径(collect_stream / nonstream)与 salvage 都以 callable
传入 —— AgentLoop 把 bound method 递进来,单测在实例上打桩即可,本模块不 import loop。
"""
from __future__ import annotations

import json
from collections.abc import Callable
from typing import Any, Dict, List, Optional, Tuple

from .storage import record_empty_response, record_malformed_tool_call

# ─────────────────────── delta / usage 提取 ───────────────────────


class PreferNonstreamToolCall(Exception):
    """流式首包已选中不适合分片传输的工具，请调用方改走非流式。"""

    def __init__(self, tool: str) -> None:
        super().__init__(tool)
        self.tool = tool


def extract_delta_content(chunk: Any) -> Optional[str]:
    """从 stream chunk 提 delta.content(文本片段)。chunk 形态 litellm ModelResponseStream:
    choices[0].delta.content。usage-only 收尾 chunk(没 choices / delta)返 None。
    """
    try:
        choices = getattr(chunk, "choices", None)
        if not choices:
            return None
        delta = getattr(choices[0], "delta", None)
        if delta is None:
            return None
        content = getattr(delta, "content", None)
        return content if content else None
    except Exception:
        return None


def extract_delta_reasoning(chunk: Any) -> Optional[str]:
    """从 stream chunk 提 delta.reasoning_content(thinking 模型的推理片段)。
    litellm 对多数 provider 归一到 delta.reasoning_content,个别只放
    provider_specific_fields —— 两处都查。没有则返 None(非 thinking 模型零开销)。
    """
    try:
        choices = getattr(chunk, "choices", None)
        if not choices:
            return None
        delta = getattr(choices[0], "delta", None)
        if delta is None:
            return None
        rc = getattr(delta, "reasoning_content", None)
        if not rc:
            psf = getattr(delta, "provider_specific_fields", None) or {}
            rc = psf.get("reasoning_content") if isinstance(psf, dict) else None
        return rc if rc else None
    except Exception:
        return None


def extract_delta_tool_names(chunk: Any) -> List[str]:
    """提取本 chunk 新出现的 tool call function name；兼容对象和 dict 形态。"""
    try:
        choices = getattr(chunk, "choices", None)
        if not choices:
            return []
        delta = getattr(choices[0], "delta", None)
        if delta is None:
            return []
        tool_calls = getattr(delta, "tool_calls", None) or []
        names: List[str] = []
        for tc in tool_calls:
            fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None)
            if fn is None:
                continue
            name = fn.get("name") if isinstance(fn, dict) else getattr(fn, "name", None)
            if name:
                names.append(str(name))
        return names
    except Exception:
        return []


def usage_to_dict(usage: Any) -> dict:
    if not usage:
        return {}
    if hasattr(usage, "model_dump"):
        usage = usage.model_dump()
    elif hasattr(usage, "dict"):
        usage = usage.dict()
    if isinstance(usage, dict):
        return usage
    return {}


def extract_usage_details(usage: Any) -> dict:
    """从 provider usage 提取统一 token 明细。

    DeepSeek 直接给 prompt_cache_hit_tokens / prompt_cache_miss_tokens;
    OpenAI 风格把 cached tokens 放在 prompt_tokens_details.cached_tokens。
    """
    data = usage_to_dict(usage)
    prompt_details = data.get("prompt_tokens_details") or {}
    completion_details = data.get("completion_tokens_details") or {}
    if not isinstance(prompt_details, dict):
        prompt_details = {}
    if not isinstance(completion_details, dict):
        completion_details = {}

    cache_hit = (
        data.get("prompt_cache_hit_tokens")
        or prompt_details.get("cached_tokens")
        or 0
    )
    cache_miss = data.get("prompt_cache_miss_tokens") or 0
    return {
        "tokens_in": int(data.get("prompt_tokens") or 0),
        "tokens_out": int(data.get("completion_tokens") or 0),
        "cache_hit_tokens": int(cache_hit or 0),
        "cache_miss_tokens": int(cache_miss or 0),
        "reasoning_tokens": int(completion_details.get("reasoning_tokens") or 0),
    }


def extract_usage(usage: Any) -> Tuple[int, int]:
    """从 litellm response.usage 提 (prompt_tokens, completion_tokens)。"""
    details = extract_usage_details(usage)
    return details["tokens_in"], details["tokens_out"]


# ─────────────────────── 故障检测 ───────────────────────

def malformed_tool_calls(response: Any) -> List[str]:
    """检出 arguments 损坏(JSON 解析不了)的 tool_call,返回 [name(len=N), ...]。

    背景:deepseek-v4-flash 大参数工具调用偶发畸形 —— 流式 delta 错位把别处的内容
    碎片粘到 arguments 开头(如 `].cells[1].merge(...{"path":...}`),拼回来后 JSON
    解析直接失败。这种是上游瞬时抖动,不该入库污染上下文,调用方据此丢弃整轮重 roll。

    只看「解析失败」;空字符串 / 合法空对象不算畸形(交给 executor 按缺参数处理)。
    """
    try:
        msg = response.choices[0].message
    except Exception:
        return []
    bad: List[str] = []
    for tc in (getattr(msg, "tool_calls", None) or []):
        raw = (getattr(tc.function, "arguments", None) or "").strip()
        if not raw:
            continue
        try:
            json.loads(raw)
        except (json.JSONDecodeError, ValueError):
            bad.append(f"{tc.function.name}(len={len(raw)})")
    return bad


def toolcalls_partial_args(
    response: Any, required_by_tool: Dict[str, List[str]]
) -> List[Tuple[Any, str, List[str]]]:
    """检出「JSON 能解析、但必填 key 被吞掉」的畸形 tool_call。

    背景(2026-07,失败面板 #3:edit `缺少必填参数 ['path']` 跨 13 task/9 用户):流式
    arguments delta 乱序把某个键的值碎片瞬移拼进相邻字符串(实证 [1]:path 的
    `.../gen_final_report_v2.py` 被吞进 old_str 尾部),独立的 `"path"` 键随之消失。这类
    与 char-0 前缀畸形的关键区别是 **JSON parse 成功** → 既不命中 `malformed_tool_calls`
    (只抓 parse 失败)、salvage 也救不了(parse-to-end/key 白名单对合法但错位无能),一路
    漏到 executor 才在语义层报「缺必填参数」,回 [Error] 喂回模型 → 再拼再乱序,反复烧。

    判据(窄,避免误伤模型真漏参):解析为**非空 dict** 且 **至少一个必填 key 在场**同时
    **至少一个必填 key 缺失**(即 0 < len(missing) < len(required))。空 `{}` / 必填全缺
    (纯垃圾/无关键)不算 —— 交给 executor + _RepeatGuard 现状处理。

    返回 [(tc, name, missing_keys), ...]。
    """
    try:
        msg = response.choices[0].message
    except Exception:
        return []
    out: List[Tuple[Any, str, List[str]]] = []
    for tc in (getattr(msg, "tool_calls", None) or []):
        try:
            name = tc.function.name
            raw = (getattr(tc.function, "arguments", None) or "").strip()
        except Exception:
            continue
        if not raw:
            continue
        try:
            obj = json.loads(raw)
        except (json.JSONDecodeError, ValueError):
            continue  # parse 失败归 malformed_tool_calls,不重复处理
        if not isinstance(obj, dict) or not obj:
            continue
        required = required_by_tool.get(name) or []
        if not required:
            continue
        missing = [k for k in required if k not in obj]
        if 0 < len(missing) < len(required):
            out.append((tc, name, missing))
    return out


def is_empty_response(response: Any) -> bool:
    """检出「空响应」:assistant 轮既无 tool_calls 又无正文(去空白后为空)。

    背景(task 2a1bc25d 案):provider wire 偶发吐空 —— 截断流 / finish_reason 无内容 /
    网关把 tool_use 漏成正文后又丢空,回来的这一轮 tool_calls 空且 content 空。run loop
    见 tool_calls 空即当「模型答完」静默 done、返回空串,无报错、run_status=idle,与卡死
    无法区分。故和畸形同类对待:丢弃本轮走非流式重试。
    注意:纯 tool_call 轮(tc 非空、content 空)不算空响应 —— 那是正常的工具调用轮。
    """
    try:
        msg = response.choices[0].message
    except (AttributeError, IndexError, TypeError):
        return False
    if getattr(msg, "tool_calls", None):
        return False
    content = getattr(msg, "content", None) or ""
    return not content.strip()


def finish_reason(response: Any) -> str:
    """取本轮 finish_reason(取不到返 "")。length=达输出上限被截断,与 wire 吐空区分开:
    截断是我方输出预算/推理失控(如 GLM thinking 烧穿),同上下文重试无效。"""
    try:
        return getattr(response.choices[0], "finish_reason", "") or ""
    except (AttributeError, IndexError, TypeError):
        return ""


# ─────────────────────── 故障留痕(stdout + usage_events 双写,静默失败)───────────────────────

def log_partial_args(
    task_id: Any, user_id: Any, model_profile: str,
    partial: List[Tuple[Any, str, List[str]]], response: Any,
) -> None:
    """必填 key 被吞的畸形留痕:与 log_malformed_args 对称,进 usage_events(kind=
    tool_malformed),error 签名固定为 `missing required keys [...]` —— 在失败面板里和
    char-0 型(`Expecting value`)、executor 的「缺必填参数」区分开,便于统计这条新裂缝。
    任何一路失败都静默,绝不打断重试主路径。"""
    try:
        usage = extract_usage_details(getattr(response, "usage", None))
        for tc, name, missing in partial:
            raw = (getattr(tc.function, "arguments", None) or "")
            err = f"missing required keys {missing}"
            print(
                f"[malformed:partial] task={task_id} tool={name} len={len(raw)} "
                f"{err} head={ascii(raw[:300])} tail={ascii(raw[-300:])}",
                flush=True,
            )
            try:
                record_malformed_tool_call(
                    task_id=task_id,
                    user_id=user_id,
                    model_profile=model_profile,
                    tool=name,
                    arg_len=len(raw),
                    error=err,
                    head=raw[:300],
                    tail=raw[-300:],
                    tokens_in=usage["tokens_in"],
                    tokens_out=usage["tokens_out"],
                )
            except Exception:
                pass
    except Exception:
        pass


def log_empty_response(
    task_id: Any, user_id: Any, model_profile: str, response: Any, attempt: int,
    finish: str = "",
) -> None:
    """空响应留痕:stdout + usage_events(kind=empty_response)双写,任何一路失败都静默。

    与 log_malformed_args 对称:空响应轮同样整轮丢弃、messages 无痕,这里是唯一留痕 ——
    供事后定性(哪个模型档在吐空)与「工具失败聚集」面板第四段聚合。留痕绝不能反过来打断
    重试主路径(单测无 DB 时 record_* 落库失败也吞掉)。
    """
    try:
        usage = extract_usage_details(getattr(response, "usage", None))
        print(
            f"[empty_response] task={task_id} mp={model_profile} attempt={attempt} "
            f"finish={finish or '?'} tok={usage['tokens_in']}/{usage['tokens_out']}",
            flush=True,
        )
        try:
            record_empty_response(
                task_id=task_id,
                user_id=user_id,
                model_profile=model_profile,
                attempt=attempt,
                tokens_in=usage["tokens_in"],
                tokens_out=usage["tokens_out"],
                finish_reason=finish,
            )
        except Exception:
            pass  # DB 不可用(如单测无 DB)不影响 stdout 留痕与重试
    except Exception:
        pass


def log_malformed_args(
    task_id: Any, user_id: Any, model_profile: str, response: Any,
) -> None:
    """畸形 arguments 的首尾片段 + JSON 报错位置留痕:stdout + usage_events 双写。

    畸形轮不 append/不记账,这里是唯一留痕 —— 用于事后定性损坏形态(provider 流式
    delta 错位 vs 本地 stream_chunk_builder 拼接 bug),以及向 provider 报 case 取证。
    stdout 片段过 ascii() 转义:日志消费端编码不可控(Windows dev 控制台 GBK 遇 emoji
    会崩),转义后 grep '\\[malformed\\]' 拿到的内容可无损还原。DB 行(kind=tool_malformed,
    cost=0)喂 admin「工具失败聚集」面板 + 巡检邮件(core/toolfail.py),units 里
    快照该轮真实 token 供估算浪费。任何一路失败都静默 —— 留痕绝不能反过来打断重试主路径。
    """
    try:
        msg = response.choices[0].message
        usage = extract_usage_details(getattr(response, "usage", None))
        for tc in (getattr(msg, "tool_calls", None) or []):
            raw = (getattr(tc.function, "arguments", None) or "").strip()
            if not raw:
                continue
            try:
                json.loads(raw)
            except (json.JSONDecodeError, ValueError) as e:
                print(
                    f"[malformed] task={task_id} tool={tc.function.name} "
                    f"len={len(raw)} err={e} "
                    f"head={ascii(raw[:300])} tail={ascii(raw[-300:])}",
                    flush=True,
                )
                try:
                    record_malformed_tool_call(
                        task_id=task_id,
                        user_id=user_id,
                        model_profile=model_profile,
                        tool=tc.function.name,
                        arg_len=len(raw),
                        error=str(e),
                        head=raw[:300],
                        tail=raw[-300:],
                        tokens_in=usage["tokens_in"],
                        tokens_out=usage["tokens_out"],
                    )
                except Exception:
                    pass  # DB 不可用(如单测无 DB)不影响 stdout 留痕与重试
    except Exception:
        pass


# ─────────────────────── 重试策略 ───────────────────────

def robust_stream(
    *,
    collect_stream: Callable[[List[dict]], Tuple[Optional[Any], bool]],
    nonstream: Callable[[List[dict]], Optional[Any]],
    try_salvage: Callable[[Any], bool],
    llm_messages: List[dict],
    required_by_tool: Dict[str, List[str]],
    emit: Callable[[dict], None],
    llm_start_event: dict,
    task_id: Any,
    user_id: Any,
    model_profile: str,
    max_attempts: int,
) -> Tuple[Optional[Any], bool]:
    """拉一轮 LLM 并保证返回的 tool_call arguments 可解析。

    返回 (response, cancelled_mid_stream):
    - 正常完结 → (response, False);response shape 与非流式 completion() 等价
    - 中途 cancel → (None, True);已收 chunk 丢弃(非流式重试期间 cancel 同样)

    畸形重试:deepseek v4 系(flash/pro 均实测踩过)大参数工具调用偶发把内容碎片
    错位粘进 arguments,拼回后 JSON 解析失败。这种损坏一旦入库会被每轮重发、诱导
    模型继续学坏(投毒级联)。故拼回后先校验 tool_call arguments 能否解析:不能 →
    丢弃整轮(不 append/不记账,原始损坏片段打服务端日志留痕)并立刻降级非流式重试
    (同轮流式失败强相关,重试不再走流式);全部尝试耗尽仍畸形则交给 executor 的
    invalid-JSON 分支返错给模型。重试消耗的 token 不单独记账。

    流式只试一次:实测(2026-07,task 716ed3be,deepseek-v4-pro)3~4k 字符中文长文
    write 的流式重 roll 同轮连挂 3 次 —— 同轮失败强相关而非独立随机,首败即降级非流式,
    历史数据里非流式兜底从未再畸形。
    """
    response = None
    for attempt in range(max_attempts):
        use_nonstream = attempt > 0
        # 每个 attempt 重发 llm_start(stats 同一份):非流式重试完成前零 delta 事件,
        # 而 warn 事件会让前端把当前文字段定稿关闭 —— 不重发的话「思考中 · Ns」占位段
        # 没人重建,页面静止到重试完成,与卡死无法区分。
        emit(dict(llm_start_event))
        if use_nonstream:
            response = nonstream(llm_messages)
            if response is None:
                # 非流式重试期间用户点了停止(线程级 poll,见 loop._nonstream_once)
                return None, True
        else:
            cancelled = False
            try:
                response, cancelled = collect_stream(llm_messages)
            except PreferNonstreamToolCall as reroute:
                # function.name 通常早于大段 arguments 到达。此时立刻关流并由 provider
                # 一次性拼好 JSON，可避开 DeepSeek 长 write/edit 参数的 delta 错位。
                emit({
                    "type": "warn",
                    "level": "info",
                    "msg": f"检测到大参数工具 {reroute.tool}，已切换稳定传输模式",
                })
                emit(dict(llm_start_event))
                response = nonstream(llm_messages)
                if response is None:
                    return None, True
                use_nonstream = True
            if cancelled:
                return None, True

        bad = malformed_tool_calls(response)
        if not bad:
            # 空响应(tc 空且正文空):provider wire 吐空,和畸形同类瞬态故障 ——
            # 丢弃本轮走非流式重试(多数瞬态重发一次即好,用户无感),留痕供面板可见。
            if is_empty_response(response):
                fr = finish_reason(response)
                log_empty_response(
                    task_id, user_id, model_profile, response, attempt + 1, finish=fr,
                )
                # length=达输出上限被截断(我方输出预算/推理烧穿,非网关 wire 吐空)——
                # 同上下文重试大概率再撞,措辞据实区分,便于用户/日志判性质(治本在
                # 模型档,如 GLM 已禁 thinking 免推理烧穿;这里保证可观测 + 不误导)。
                truncated = fr == "length"
                emit({
                    "type": "warn",
                    "msg": (
                        ("模型输出达上限被截断" if truncated else "模型返回空响应")
                        + ",丢弃本轮"
                        f"{'重试' if use_nonstream else ',改非流式重试'}"
                        f" ({attempt + 1}/{max_attempts})"
                    ),
                })
                continue
            # 必填 key 被吞的畸形(parse 成功、salvage 无能):非流式重试(服务端一次拼好,
            # 绕开流式 delta 乱序)。耗尽尝试仍缺 → 落下面 return,交 executor 返「缺必填参数」
            # 给模型(多为模型真漏参,不再空转)。
            partial = toolcalls_partial_args(response, required_by_tool)
            if partial:
                log_partial_args(task_id, user_id, model_profile, partial, response)
                names = ", ".join(f"{n}(missing={m})" for _, n, m in partial)
                emit({
                    "type": "warn",
                    "msg": (
                        f"工具调用必填参数被吞 {names},丢弃本轮"
                        f"{'重试' if use_nonstream else ',改非流式重试'}"
                        f" ({attempt + 1}/{max_attempts})"
                    ),
                })
                continue
            return response, False
        # 先尝试就地抢救:畸形是 char-0 垃圾前缀 + 尾部完好 JSON(定层已证 provider-wire),
        # 全部畸形 tool_call 都能抠出干净 JSON 才改写并当轮继续,省掉一次非流式重试;
        # 任一抠不出则一个都不动,原样走下面的丢弃 + 重试(零回退风险)。
        if try_salvage(response):
            return response, False
        log_malformed_args(task_id, user_id, model_profile, response)
        emit({
            "type": "warn",
            "msg": (
                f"工具调用参数损坏 {bad},丢弃本轮"
                f"{'重试' if use_nonstream else ',改非流式重试'}"
                f" ({attempt + 1}/{max_attempts})"
            ),
        })
    # 非流式重试仍畸形(理论极罕见):交还给 _execute_tool_call 的 invalid-JSON 分支
    # 优雅返错给模型,而非在此死循环。
    return response, False
