"""
Streaming probe for OpenAI-compatible backends: detect whether ``delta.tool_calls``
is populated (native tool-call path) vs. tool markup only in ``content`` (XML/Hermes fallback).

Used from preferences (MLX) to catch misconfigured servers or missing tool-call parsers.
"""

from __future__ import annotations

import json
import logging
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

import httpx
from openai import OpenAI

_log = logging.getLogger(__name__)

PROBE_TOOL_NAME = "mugi_native_tool_probe"

_PROBE_TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": PROBE_TOOL_NAME,
            "description": (
                "Connectivity probe only. When the user asks to run the native tool probe, "
                "you must call this function with ping set to the string ok."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "ping": {
                        "type": "string",
                        "description": "Echo string; use ok for the probe.",
                    }
                },
                "required": [],
            },
        },
    }
]

_USER_PROBE_MESSAGE = (
    "Run the native tool probe: you MUST call the function "
    f"{PROBE_TOOL_NAME} with ping set to the string ok. Do not answer in plain text only."
)


def _normalize_v1_base_url(base_url: str) -> str:
    bu = (base_url or "").strip()
    if not bu:
        return "http://127.0.0.1:8000/v1"
    if not bu.endswith("/v1"):
        bu = bu.rstrip("/") + "/v1"
    return bu


def _fetch_first_model_id(base_url: str, timeout: float = 15.0) -> str | None:
    """Resolve the model id to probe, preferring the *actually loaded* model.

    CRITICAL: mlx-vlm's ``GET /v1/models`` lists the entire HF cache, NOT just the loaded
    model — ``data[0]`` is often some unrelated small model. The server serves a single model
    at a time, so probing ``data[0]`` would EVICT the user's loaded chat model and load the
    probe's pick (we saw a 26B chat model evicted by a 0.5B from the probe). So we read
    ``GET /health.loaded_model`` first and only fall back to ``/v1/models`` ``data[0]`` when
    /health is unavailable (e.g. non-mlx-vlm OpenAI servers).
    """
    bu = base_url.rstrip("/")
    root = bu[:-3].rstrip("/") if bu.endswith("/v1") else bu
    try:
        with urlopen(root + "/health", timeout=timeout) as resp:
            health = json.loads(resp.read().decode("utf-8", errors="replace"))
        loaded = str((health or {}).get("loaded_model") or "").strip() if isinstance(health, dict) else ""
        if loaded:
            return loaded
    except (HTTPError, URLError, OSError, json.JSONDecodeError) as e:
        _log.debug("probe: /health unavailable, falling back to /v1/models: %s", e)
    try:
        with urlopen(bu + "/models", timeout=timeout) as resp:
            payload = json.loads(resp.read().decode("utf-8", errors="replace"))
    except (HTTPError, URLError, OSError, json.JSONDecodeError) as e:
        _log.debug("probe: models list failed: %s", e)
        return None
    data = payload.get("data") if isinstance(payload, dict) else None
    if not isinstance(data, list) or not data:
        return None
    mid = str((data[0] or {}).get("id") or "").strip()
    return mid or None


def _delta_tool_call_has_name(delta: Any) -> bool:
    """True if this delta includes any tool call fragment with a non-empty function name."""
    tcs = getattr(delta, "tool_calls", None) or []
    for tc in tcs:
        if isinstance(tc, dict):
            fn = tc.get("function")
            name = (fn.get("name") if isinstance(fn, dict) else None) or ""
            if str(name).strip():
                return True
            continue
        fn = getattr(tc, "function", None)
        if fn is not None:
            name = getattr(fn, "name", None) or ""
            if str(name).strip():
                return True
    return False


def _delta_has_native_tool_call_signal(delta: Any) -> bool:
    """True if streaming deltas show OpenAI-style tool_calls (not only Hermes/XML in content).

    Some OpenAI-compatible servers emit ``delta.tool_calls`` with ``id`` + ``type=function``
    before the ``function.name`` fragment arrives; counting only names caused false negatives.
    """
    if _delta_tool_call_has_name(delta):
        return True
    tcs = getattr(delta, "tool_calls", None) or []
    for tc in tcs:
        if isinstance(tc, dict):
            tc_id = str(tc.get("id") or "").strip()
            tc_type = str(tc.get("type") or "").strip().lower()
            if tc_id and tc_type == "function":
                return True
            fn = tc.get("function")
            if isinstance(fn, dict):
                args = fn.get("arguments")
                if args is not None and str(args).strip():
                    return True
            continue
        tc_id = str(getattr(tc, "id", None) or "").strip()
        tc_type = str(getattr(tc, "type", None) or "").strip().lower()
        if tc_id and tc_type == "function":
            return True
        fn = getattr(tc, "function", None)
        if fn is not None:
            args = getattr(fn, "arguments", None)
            if args is not None and str(args).strip():
                return True
    return False


def _content_has_xml_tool_markup(content: str) -> bool:
    lowered = (content or "").lower()
    return "<tool_call" in lowered or "<|tool_call" in lowered


def _message_has_native_tool_calls(message: Any) -> bool:
    """True if a chat completion ``message`` includes OpenAI-style ``tool_calls``."""
    tcs = getattr(message, "tool_calls", None) or []
    return bool(tcs)


def _probe_stream_once(
    client: OpenAI,
    *,
    resolved_model: str,
    tool_choice: Any,
    max_tokens: int,
) -> tuple[bool, str, str | None]:
    """One streaming attempt; returns (native_detected, concatenated_content, stream_error_or_none)."""
    native = False
    content_buf = ""
    stream_err: str | None = None
    try:
        stream = client.chat.completions.create(
            model=resolved_model,
            messages=[{"role": "user", "content": _USER_PROBE_MESSAGE}],
            tools=_PROBE_TOOLS,
            tool_choice=tool_choice,
            stream=True,
            max_tokens=max_tokens,
            temperature=0.0,
            frequency_penalty=0.0,
            presence_penalty=0.0,
        )
        for chunk in stream:
            if not chunk.choices:
                continue
            choice = chunk.choices[0]
            delta = choice.delta
            if _delta_has_native_tool_call_signal(delta):
                native = True
            # Some OpenAI-compatible servers omit incremental
            # ``delta.tool_calls`` fragments but set ``finish_reason`` to ``tool_calls`` on the last chunk.
            fr = getattr(choice, "finish_reason", None)
            if fr == "tool_calls":
                native = True
            c = getattr(delta, "content", None)
            if c:
                content_buf += c
    except Exception as e:
        stream_err = str(e)
        _log.debug("probe: stream failed (tool_choice=%r): %s", tool_choice, e, exc_info=True)
    return native, content_buf, stream_err


def probe_streaming_native_tool_calls(
    base_url: str,
    *,
    api_key: str = "lm-studio",
    model: str | None = None,
    max_tokens: int = 128,
    timeout: float = 60.0,
) -> dict[str, Any]:
    """
    One streaming chat completion with a single dummy tool; checks whether the server
    emits native ``delta.tool_calls``.

    Returns a JSON-serializable dict with keys including:
    - ``ok``: True if streaming carried native tool signals, or a non-streaming completion returned
      ``message.tool_calls`` (some OpenAI-compatible servers).
    - ``non_stream_fallback``: True when ``ok`` was established only via the non-streaming check.
    - ``native_tool_calls``: same as ``ok`` (explicit).
    - ``xml_like_in_content``: assistant ``content`` contained ``<tool_call``-style markup.
    - ``required_tool_choice_retry``: True when a second pass used ``tool_choice="required"``.
    - ``model``, ``base_url``: what was used.
    - ``message``: short human-readable summary.
    - ``error``: set on transport/SDK failure (``ok`` will be False).
    """
    bu = _normalize_v1_base_url(base_url)
    resolved_model = (model or "").strip() or None
    if not resolved_model:
        resolved_model = _fetch_first_model_id(bu, timeout=min(15.0, timeout))
    if not resolved_model:
        return {
            "ok": False,
            "native_tool_calls": False,
            "non_stream_fallback": False,
            "xml_like_in_content": False,
            "required_tool_choice_retry": False,
            "model": "",
            "base_url": bu,
            "message": "No model id available (GET /v1/models returned empty).",
            "error_code": "no_model",
            "error_detail": None,
        }

    client = OpenAI(
        base_url=bu,
        api_key=api_key or "lm-studio",
        timeout=httpx.Timeout(connect=30.0, read=float(timeout), write=30.0, pool=30.0),
        max_retries=0,
    )

    mt = max(16, min(int(max_tokens), 512))
    native, content_buf, stream_err = _probe_stream_once(
        client,
        resolved_model=resolved_model,
        tool_choice="auto",
        max_tokens=mt,
    )
    required_retry = False
    if not native and stream_err is None:
        required_retry = True
        native2, buf2, stream_err2 = _probe_stream_once(
            client,
            resolved_model=resolved_model,
            tool_choice="required",
            max_tokens=mt,
        )
        native = native or native2
        content_buf += buf2
        if stream_err2:
            _log.debug(
                "probe: tool_choice=required pass failed (ignored unless primary failed): %s",
                stream_err2[:400],
            )

    native_from_stream = bool(native)
    non_stream_native = False
    if not native_from_stream and stream_err is None:
        try:
            resp = client.chat.completions.create(
                model=resolved_model,
                messages=[{"role": "user", "content": _USER_PROBE_MESSAGE}],
                tools=_PROBE_TOOLS,
                tool_choice="required",
                stream=False,
                max_tokens=mt,
                temperature=0.0,
                frequency_penalty=0.0,
                presence_penalty=0.0,
            )
            ch0 = resp.choices[0] if getattr(resp, "choices", None) else None
            msg = getattr(ch0, "message", None) if ch0 else None
            if _message_has_native_tool_calls(msg):
                non_stream_native = True
        except Exception as e:
            _log.debug("probe: non-stream completion failed: %s", e, exc_info=True)

    native = native_from_stream or non_stream_native
    non_stream_fallback = bool(non_stream_native and not native_from_stream)

    xml_like = _content_has_xml_tool_markup(content_buf)
    if native:
        if non_stream_fallback:
            msg = (
                "Native tool_calls confirmed via non-streaming completion; streaming did not expose "
                "delta.tool_calls fragments. Chat uses streaming — update the mlx-vlm server if tools fail in threads."
            )
            code = None
            detail = None
        elif required_retry and native_from_stream:
            msg = (
                "Streaming native tool_calls detected after tool_choice=required retry — "
                "first pass (auto) did not emit delta.tool_calls."
            )
            code = None
            detail = None
        else:
            msg = "Streaming native tool_calls detected — server is populating delta.tool_calls."
            code = None
            detail = None
    elif stream_err:
        msg = f"Probe request failed: {stream_err[:400]}"
        code = "stream_error"
        detail = stream_err[:800]
    elif xml_like:
        msg = (
            "No native streaming tool_calls; assistant content looks like XML/Hermes tool markup. "
            "The mlx-vlm server auto-infers the tool parser from the model's chat template — try a "
            "different model / restart the server after a model change."
        )
        code = "xml_only"
        detail = None
    else:
        msg = (
            "No native streaming tool_calls and no XML tool markup detected — "
            "model may have refused the probe or returned plain text."
        )
        code = "no_tool_signal"
        detail = None

    return {
        "ok": native,
        "native_tool_calls": native,
        "non_stream_fallback": non_stream_fallback,
        "xml_like_in_content": xml_like,
        "required_tool_choice_retry": required_retry,
        "model": resolved_model,
        "base_url": bu,
        "message": msg,
        "error_code": code,
        "error_detail": detail,
    }
