"""Normalize OpenAI-compatible thinking fields across provider versions.

Older vLLM / DeepSeek put thinking in ``reasoning_content``. Newer vLLM
reasoning parsers put it in ``reasoning``. DeerFlow's UI and downstream code
canonicalize on ``AIMessage.additional_kwargs.reasoning_content``.

Only fields that actually appeared on a response are echoed back on later
turns, so models that still use ``reasoning_content`` (or send neither) are
unchanged.
"""

from __future__ import annotations

import json
from typing import Any

from langchain_core.messages import AIMessage, AIMessageChunk


def reasoning_to_text(reasoning: Any) -> str:
    """Best-effort extraction of readable reasoning text from provider payloads."""
    if isinstance(reasoning, str):
        return reasoning

    if isinstance(reasoning, list):
        parts = [reasoning_to_text(item) for item in reasoning]
        return "".join(part for part in parts if part)

    if isinstance(reasoning, dict):
        for key in ("text", "content", "reasoning"):
            value = reasoning.get(key)
            if isinstance(value, str):
                return value
            if value is not None:
                text = reasoning_to_text(value)
                if text:
                    return text
        try:
            return json.dumps(reasoning, ensure_ascii=False)
        except TypeError:
            return str(reasoning)

    if reasoning is None:
        return ""
    try:
        return json.dumps(reasoning, ensure_ascii=False)
    except TypeError:
        return str(reasoning)


def extract_reasoning_text(payload: dict | None) -> str | None:
    """Read thinking text from ``reasoning_content`` or ``reasoning``.

    Prefer ``reasoning_content`` when both are present so older DeepSeek /
    vLLM payloads keep their original string. Newer vLLM that only sends
    ``reasoning`` (string or structured) is still captured.
    """
    if not isinstance(payload, dict):
        return None
    for key in ("reasoning_content", "reasoning"):
        value = payload.get(key)
        if value is None:
            continue
        text = reasoning_to_text(value)
        if text:
            return text
    return None


def _merge_streamed_value(existing: Any, incoming: Any) -> Any:
    if isinstance(existing, str) and isinstance(incoming, str):
        if incoming and incoming not in existing:
            return f"{existing}{incoming}"
        return existing or incoming
    return incoming if incoming is not None else existing


def apply_reasoning_from_payload(
    message: AIMessage | AIMessageChunk,
    payload: dict | None,
    *,
    append: bool = True,
) -> AIMessage | AIMessageChunk:
    """Copy provider thinking fields onto ``additional_kwargs`` for UI + echo-back."""
    if not isinstance(payload, dict):
        return message

    additional_kwargs = dict(message.additional_kwargs)
    changed = False

    if "reasoning" in payload and payload.get("reasoning") is not None:
        incoming = payload.get("reasoning")
        existing = additional_kwargs.get("reasoning")
        merged = _merge_streamed_value(existing, incoming) if append else incoming
        if merged is not None and merged != existing:
            additional_kwargs["reasoning"] = merged
            changed = True

    text = extract_reasoning_text(payload)
    if text:
        existing = additional_kwargs.get("reasoning_content")
        if append and isinstance(existing, str) and existing:
            if text not in existing:
                additional_kwargs["reasoning_content"] = f"{existing}{text}"
                changed = True
        elif text != existing:
            additional_kwargs["reasoning_content"] = text
            changed = True

    if not changed:
        return message
    return message.model_copy(update={"additional_kwargs": additional_kwargs})


def fill_ui_reasoning_content(message: AIMessage | AIMessageChunk) -> AIMessage | AIMessageChunk:
    """If LangChain stored ``reasoning`` but not the UI field, copy the text across."""
    kwargs = message.additional_kwargs
    reasoning = kwargs.get("reasoning")
    if reasoning is None or kwargs.get("reasoning_content"):
        return message
    return apply_reasoning_from_payload(message, {"reasoning": reasoning}, append=False)


def restore_reasoning_fields(payload_msg: dict, orig_msg: AIMessage) -> None:
    """Re-inject the thinking field this assistant turn originally carried.

    Newer vLLM stores thinking on ``reasoning`` and we also copy a UI-facing
    ``reasoning_content`` string from it. Echo the wire field that actually
    arrived: ``reasoning`` wins when present so we do not send a synthesized
    ``reasoning_content`` to a parser that no longer accepts it. Older
    DeepSeek / vLLM turns only have ``reasoning_content`` and stay on that
    field.
    """
    reasoning = orig_msg.additional_kwargs.get("reasoning")
    if reasoning is not None:
        payload_msg["reasoning"] = reasoning
        return
    reasoning_content = orig_msg.additional_kwargs.get("reasoning_content")
    if reasoning_content is not None:
        payload_msg["reasoning_content"] = reasoning_content