import logging

from langchain.chat_models import BaseChatModel

from deerflow.config import get_app_config
from deerflow.config.app_config import AppConfig
from deerflow.reflection import resolve_class
from deerflow.tracing import build_tracing_callbacks

logger = logging.getLogger(__name__)


def _deep_merge_dicts(base: dict | None, override: dict) -> dict:
    """Recursively merge two dictionaries without mutating the inputs."""
    merged = dict(base or {})
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = _deep_merge_dicts(merged[key], value)
        else:
            merged[key] = value
    return merged


def _vllm_disable_chat_template_kwargs(chat_template_kwargs: dict) -> dict:
    """Build the disable payload for vLLM/Qwen chat template kwargs."""
    disable_kwargs: dict[str, bool] = {}
    if "thinking" in chat_template_kwargs:
        disable_kwargs["thinking"] = False
    if "enable_thinking" in chat_template_kwargs:
        disable_kwargs["enable_thinking"] = False
    return disable_kwargs


# Qwen3.8 chat templates only accept these three effort names.
_QWEN_REASONING_EFFORT = {
    "minimal": "low",
    "low": "low",
    "medium": "medium",
    "high": "xhigh",
    "xhigh": "xhigh",
}

_OPENAI_THINKING_OFF_EXTRA_BODY = {
    "enable_thinking": False,
    "chat_template_kwargs": {"enable_thinking": False},
    "thinking": {"type": "disabled"},
}


def _is_vllm_provider(model_use: str) -> bool:
    lowered = (model_use or "").lower()
    return "vllm_provider" in lowered or lowered.endswith(":vllmchatmodel")


def _is_qwen38_model(model_config) -> bool:
    blob = f"{getattr(model_config, 'name', '')} {getattr(model_config, 'model', '')}".lower()
    return "qwen3.8" in blob


def _should_control_openai_thinking(model_config) -> bool:
    """Whether this OpenAI-compatible model needs explicit thinking on/off flags.

    Qwen3.8 (and current vLLM reasoning parsers) default to thinking at ``xhigh``.
    If DeerFlow does not send ``enable_thinking: false``, the first visible answer
    token waits until that hidden reasoning finishes — third-party chat UIs that
    disable thinking or stream the reasoning field therefore feel much faster.
    """
    if getattr(model_config, "supports_thinking", False):
        return True
    if _is_vllm_provider(getattr(model_config, "use", "")):
        return True
    return _is_qwen38_model(model_config)


def _map_qwen_reasoning_effort(effort: object) -> str | None:
    if effort is None:
        return None
    return _QWEN_REASONING_EFFORT.get(str(effort).strip().lower())


def _apply_openai_thinking_control(
    model_config,
    model_class,
    model_settings_from_config: dict,
    *,
    thinking_enabled: bool,
    reasoning_effort: object,
) -> None:
    """Inject vLLM/Qwen thinking flags when the UI toggle is on or off."""
    if "extra_body" not in getattr(model_class, "model_fields", {}):
        return
    qwen_vllm = _is_vllm_provider(getattr(model_config, "use", "")) or _is_qwen38_model(model_config)
    if thinking_enabled:
        if not qwen_vllm:
            return
        enable_body: dict = {
            "enable_thinking": True,
            "chat_template_kwargs": {"enable_thinking": True},
        }
        mapped_effort = _map_qwen_reasoning_effort(reasoning_effort)
        if mapped_effort:
            enable_body["chat_template_kwargs"]["reasoning_effort"] = mapped_effort
        model_settings_from_config["extra_body"] = _deep_merge_dicts(
            model_settings_from_config.get("extra_body"),
            enable_body,
        )
        return
    if not _should_control_openai_thinking(model_config):
        return
    disable_body = (
        {
            "enable_thinking": False,
            "chat_template_kwargs": {"enable_thinking": False},
        }
        if qwen_vllm
        else _OPENAI_THINKING_OFF_EXTRA_BODY
    )
    model_settings_from_config["extra_body"] = _deep_merge_dicts(
        model_settings_from_config.get("extra_body"),
        disable_body,
    )


def _enable_stream_usage_by_default(model_use_path: str, model_settings_from_config: dict) -> None:
    """Enable stream usage for OpenAI-compatible models unless explicitly configured.

    LangChain only auto-enables ``stream_usage`` for OpenAI models when no custom
    base URL or client is configured. DeerFlow frequently uses OpenAI-compatible
    gateways, so token usage tracking would otherwise stay empty and the
    TokenUsageMiddleware would have nothing to log.
    """
    if model_use_path != "langchain_openai:ChatOpenAI":
        return
    if "stream_usage" in model_settings_from_config:
        return
    if "base_url" in model_settings_from_config or "openai_api_base" in model_settings_from_config:
        model_settings_from_config["stream_usage"] = True


def create_chat_model(name: str | None = None, thinking_enabled: bool = False, *, app_config: AppConfig | None = None, force_disable_thinking: bool = False, **kwargs) -> BaseChatModel:
    """Create a chat model instance from the config.

    Args:
        name: The name of the model to create. If None, the first model in the config will be used.
        thinking_enabled: Turn the model's extended thinking / reasoning on.
        force_disable_thinking: Hard override that disables thinking even when the
            per-model config declares no ``when_thinking_*`` shape. Many internal
            vLLM / Qwen models only set ``supports_thinking: true`` (no disable
            shape), so passing ``thinking_enabled=False`` alone injects nothing and
            the model keeps reasoning on. When this flag is set (and thinking is
            off) we additionally inject both common OpenAI-compatible disable shapes
            — vLLM/Qwen ``chat_template_kwargs.enable_thinking=False`` and Zhipu GLM
            ``thinking.type="disabled"`` — mirroring
            ``TitleMiddleware._force_disable_thinking``. No effect when thinking is
            enabled, and skipped for model classes that don't accept ``extra_body``.

    Returns:
        A chat model instance.
    """
    config = app_config or get_app_config()
    if name is None:
        name = config.models[0].name
    model_config = config.get_model_config(name)
    if model_config is None:
        raise ValueError(f"Model {name} not found in config") from None
    model_class = resolve_class(model_config.use, BaseChatModel)
    # Transparently apply DeerFlow's OpenAI-compatible patch layer so users
    # don't need to switch every `use: langchain_openai:ChatOpenAI` entry.
    if model_config.use == "langchain_openai:ChatOpenAI":
        from deerflow.models.patched_openai import PatchedChatOpenAI

        model_class = PatchedChatOpenAI
    model_settings_from_config = model_config.model_dump(
        exclude_none=True,
        exclude={
            "use",
            "name",
            "display_name",
            "description",
            "supports_thinking",
            "supports_reasoning_effort",
            "when_thinking_enabled",
            "when_thinking_disabled",
            "thinking",
            "supports_vision",
        },
    )
    # ``max_input_tokens`` is a DeerFlow-side safety limit, not an OpenAI API
    # request parameter.  ModelConfig intentionally allows provider-specific
    # extras, so without removing it here ChatOpenAI would forward the unknown
    # key to some compatible gateways and turn the protection itself into a
    # HTTP 400.  Mirror it into LangChain's local-only ``profile`` metadata
    # when the selected model supports that field; the final context guard
    # reads the original model config and therefore also works for adapters
    # without profiles.
    configured_max_input_tokens = model_settings_from_config.pop("max_input_tokens", None)
    if configured_max_input_tokens is not None and "profile" in getattr(model_class, "model_fields", {}):
        profile = dict(model_settings_from_config.get("profile") or {})
        profile["max_input_tokens"] = configured_max_input_tokens
        model_settings_from_config["profile"] = profile
    # Compute effective when_thinking_enabled by merging in the `thinking` shortcut field.
    # The `thinking` shortcut is equivalent to setting when_thinking_enabled["thinking"].
    has_thinking_settings = (model_config.when_thinking_enabled is not None) or (model_config.thinking is not None)
    effective_wte: dict = dict(model_config.when_thinking_enabled) if model_config.when_thinking_enabled else {}
    if model_config.thinking is not None:
        merged_thinking = {**(effective_wte.get("thinking") or {}), **model_config.thinking}
        effective_wte = {**effective_wte, "thinking": merged_thinking}
    if thinking_enabled and has_thinking_settings:
        if not model_config.supports_thinking:
            raise ValueError(f"Model {name} does not support thinking. Set `supports_thinking` to true in the `config.yaml` to enable thinking.") from None
        if effective_wte:
            model_settings_from_config.update(effective_wte)
    elif thinking_enabled and model_config.supports_thinking and model_config.use == "langchain_openai:ChatOpenAI":
        # vLLM/Qwen OpenAI-compatible endpoints commonly require this explicit
        # request flag before they emit delta.reasoning.
        model_settings_from_config["extra_body"] = _deep_merge_dicts(
            model_settings_from_config.get("extra_body"),
            {"enable_thinking": True},
        )
    if not thinking_enabled:
        if model_config.when_thinking_disabled is not None:
            # User-provided disable settings take full precedence
            model_settings_from_config.update(model_config.when_thinking_disabled)
        elif has_thinking_settings and effective_wte.get("extra_body", {}).get("thinking", {}).get("type"):
            # OpenAI-compatible gateway: thinking is nested under extra_body
            model_settings_from_config["extra_body"] = _deep_merge_dicts(
                model_settings_from_config.get("extra_body"),
                {"thinking": {"type": "disabled"}},
            )
            model_settings_from_config["reasoning_effort"] = "minimal"
        elif has_thinking_settings and (disable_chat_template_kwargs := _vllm_disable_chat_template_kwargs(effective_wte.get("extra_body", {}).get("chat_template_kwargs") or {})):
            # vLLM uses chat template kwargs to switch thinking on/off.
            model_settings_from_config["extra_body"] = _deep_merge_dicts(
                model_settings_from_config.get("extra_body"),
                {"chat_template_kwargs": disable_chat_template_kwargs},
            )
        elif has_thinking_settings and effective_wte.get("thinking", {}).get("type"):
            # Native langchain_anthropic: thinking is a direct constructor parameter
            model_settings_from_config["thinking"] = {"type": "disabled"}
    requested_reasoning_effort = kwargs.get("reasoning_effort")
    if not model_config.supports_reasoning_effort:
        kwargs.pop("reasoning_effort", None)
        model_settings_from_config.pop("reasoning_effort", None)

    _apply_openai_thinking_control(
        model_config,
        model_class,
        model_settings_from_config,
        thinking_enabled=thinking_enabled,
        reasoning_effort=requested_reasoning_effort,
    )

    # Hard override: force-disable thinking regardless of the per-model config shape.
    # Injected only for OpenAI-compatible classes (those declaring `extra_body`);
    # other providers would reject the kwarg, so they are skipped (they ignore it).
    # Three disable shapes are injected so the common OpenAI-compatible gateways are
    # all covered: SiliconFlow / DashScope-style **top-level** ``enable_thinking``
    # (DeepSeek-V3.1+/Qwen3/GLM hybrid models on api.siliconflow.cn — the knob the
    # nested ``chat_template_kwargs`` shape does NOT reach), vLLM/Qwen native
    # ``chat_template_kwargs.enable_thinking``, and Zhipu GLM ``thinking.type``.
    # Only models that explicitly opt into thinking should receive vendor-specific
    # disable payloads. Generic OpenAI-compatible gateways (including LiteLLM)
    # may reject these otherwise-unknown fields instead of forwarding them.
    if (
        force_disable_thinking
        and not thinking_enabled
        and model_config.supports_thinking
        and "extra_body" in getattr(model_class, "model_fields", {})
    ):
        model_settings_from_config["extra_body"] = _deep_merge_dicts(
            model_settings_from_config.get("extra_body"),
            {"enable_thinking": False, "chat_template_kwargs": {"enable_thinking": False}, "thinking": {"type": "disabled"}},
        )

    _enable_stream_usage_by_default(model_config.use, model_settings_from_config)

    # For Codex Responses API models: map thinking mode to reasoning_effort
    from deerflow.models.openai_codex_provider import CodexChatModel

    if issubclass(model_class, CodexChatModel):
        # The ChatGPT Codex endpoint currently rejects max_tokens/max_output_tokens.
        model_settings_from_config.pop("max_tokens", None)

        # Use explicit reasoning_effort from frontend if provided (low/medium/high)
        explicit_effort = kwargs.pop("reasoning_effort", None)
        if not thinking_enabled:
            model_settings_from_config["reasoning_effort"] = "none"
        elif explicit_effort and explicit_effort in ("low", "medium", "high", "xhigh"):
            model_settings_from_config["reasoning_effort"] = explicit_effort
        elif "reasoning_effort" not in model_settings_from_config:
            model_settings_from_config["reasoning_effort"] = "medium"

    # For MindIE models: enforce conservative retry defaults.
    # Timeout normalization is handled inside MindIEChatModel itself.
    if getattr(model_class, "__name__", "") == "MindIEChatModel":
        # Enforce max_retries constraint to prevent cascading timeouts.
        model_settings_from_config["max_retries"] = model_settings_from_config.get("max_retries", 1)

    # Ensure stream_usage is enabled so that token usage metadata is available
    # in streaming responses.  LangChain's BaseChatOpenAI only defaults
    # stream_usage=True when no custom base_url/api_base is set, so models
    # hitting third-party endpoints (e.g. doubao, deepseek) silently lose
    # usage data.  We default it to True unless explicitly configured.
    if "stream_usage" not in model_settings_from_config and "stream_usage" not in kwargs:
        if "stream_usage" in getattr(model_class, "model_fields", {}):
            model_settings_from_config["stream_usage"] = True

    model_instance = model_class(**kwargs, **model_settings_from_config)

    # Always emit a compact model lifecycle trace.  In particular, the first
    # streamed callback proves that the upstream endpoint accepted the request
    # and started returning data, including reasoning-only chunks whose normal
    # content string is empty.
    from deerflow.models.timing_callback import ModelTimingCallback

    existing_callbacks = model_instance.callbacks or []
    model_instance.callbacks = [
        *existing_callbacks,
        ModelTimingCallback(
            model_name=name,
            base_url=str(model_config.base_url) if model_config.base_url else None,
        ),
    ]

    callbacks = build_tracing_callbacks()
    if callbacks:
        existing_callbacks = model_instance.callbacks or []
        model_instance.callbacks = [*existing_callbacks, *callbacks]
        logger.debug(f"Tracing attached to model '{name}' with providers={len(callbacks)}")
    return model_instance