"""End-to-end reasoning parser tests against a running vLLM server.

Tests all three thinking modes (enabled / disabled / adaptive) across
streaming and non-streaming chat completions.

Usage:
    python /tmp/test_reasoning_endpoints.py
"""

import json
import sys
import time
import urllib.request
from typing import Any

ENDPOINT = "http://localhost:8000/v1/chat/completions"
MODEL = "coder"
PROMPT = "用一句话介绍你自己"  # short, deterministic enough for a smoke test


def post(payload: dict[str, Any], stream: bool) -> dict[str, Any] | list[dict[str, Any]]:
    """Send one request. If stream=True, return list of SSE chunks."""
    body = json.dumps({**payload, "stream": stream}).encode()
    req = urllib.request.Request(
        ENDPOINT,
        data=body,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        if not stream:
            return json.loads(resp.read())
        chunks: list[dict[str, Any]] = []
        for raw in resp:
            line = raw.decode().strip()
            if not line.startswith("data:"):
                continue
            data = line[len("data:"):].strip()
            if data == "[DONE]":
                break
            chunks.append(json.loads(data))
        return chunks


def check_non_stream(label: str, payload: dict[str, Any]) -> bool:
    print(f"\n--- {label} (non-stream) ---")
    payload = {**payload, "stream": False, "max_tokens": 200, "temperature": 0.0}
    t0 = time.time()
    res = post(payload, stream=False)
    elapsed = time.time() - t0
    choice = res["choices"][0]
    msg = choice["message"]
    reasoning = msg.get("reasoning") or ""
    content = msg.get("content") or ""
    print(f"  finish_reason={choice.get('finish_reason')!r}")
    print(f"  reasoning: {reasoning!r}")
    print(f"  content:   {content!r}")
    print(f"  latency:   {elapsed:.2f}s")

    issues: list[str] = []
    mode = payload.get("chat_template_kwargs", {}).get("thinking_mode", "adaptive")
    if mode == "enabled":
        # Reasoning should be present and start with <think>-ish text;
        # content should NOT contain raw <think> or </mm:think> tags.
        if not reasoning:
            issues.append("enabled mode: empty reasoning")
        if "<mm:think>" in content or "</mm:think>" in content:
            issues.append(f"enabled mode: raw marker leaked into content: {content!r}")
    elif mode == "disabled":
        # Reasoning should be empty, content should be the direct answer
        # (no <mm:think> prefix).
        if reasoning:
            issues.append(f"disabled mode: leaked reasoning: {reasoning!r}")
        if not content:
            issues.append("disabled mode: empty content")
        if content.startswith("<mm:think>"):
            issues.append(f"disabled mode: content starts with <think>: {content!r}")
    else:  # adaptive
        # Either path is fine; just ensure raw markers don't leak into content.
        if "<mm:think>" in content or "</mm:think>" in content:
            issues.append(f"adaptive mode: raw marker leaked into content: {content!r}")
    if issues:
        for i in issues:
            print(f"  [ISSUE] {i}")
        return False
    print("  [OK]")
    return True


def check_stream(label: str, payload: dict[str, Any]) -> bool:
    print(f"\n--- {label} (stream) ---")
    payload = {**payload, "stream": True, "max_tokens": 200, "temperature": 0.0}
    t0 = time.time()
    chunks = post(payload, stream=True)
    elapsed = time.time() - t0

    # Reassemble reasoning + content from deltas.
    reasoning_parts: list[str] = []
    content_parts: list[str] = []
    finish_reason = None
    saw_reasoning_delta = False
    saw_content_delta = False
    for c in chunks:
        ch = c.get("choices", [{}])[0]
        delta = ch.get("delta", {})
        if "reasoning" in delta and delta["reasoning"] is not None:
            reasoning_parts.append(delta["reasoning"])
            saw_reasoning_delta = True
        if "content" in delta and delta["content"] is not None:
            content_parts.append(delta["content"])
            saw_content_delta = True
        if ch.get("finish_reason"):
            finish_reason = ch["finish_reason"]

    reasoning = "".join(reasoning_parts)
    content = "".join(content_parts)
    print(f"  chunks: {len(chunks)}, finish_reason={finish_reason!r}")
    print(f"  reasoning ({len(reasoning)} chars): {reasoning!r}")
    print(f"  content   ({len(content)} chars): {content!r}")
    print(f"  saw reasoning_delta={saw_reasoning_delta}, content_delta={saw_content_delta}")
    print(f"  latency: {elapsed:.2f}s")

    issues: list[str] = []
    mode = payload.get("chat_template_kwargs", {}).get("thinking_mode", "adaptive")
    if mode == "enabled":
        if not reasoning:
            issues.append("enabled mode: empty reasoning across all deltas")
        if "<mm:think>" in content or "</mm:think>" in content:
            issues.append(f"enabled mode: raw marker leaked into streamed content: {content!r}")
    elif mode == "disabled":
        if reasoning:
            issues.append(f"disabled mode: reasoning streamed: {reasoning!r}")
        if not content:
            issues.append("disabled mode: empty streamed content")
        if content.startswith("<mm:think>"):
            issues.append(f"disabled mode: streamed content starts with <think>: {content!r}")
    else:
        if "<mm:think>" in content or "</mm:think>" in content:
            issues.append(f"adaptive mode: raw marker leaked into streamed content: {content!r}")
    if issues:
        for i in issues:
            print(f"  [ISSUE] {i}")
        return False
    print("  [OK]")
    return True


def main() -> None:
    base_payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": PROMPT}],
    }

    results: list[bool] = []

    for mode in ["enabled", "disabled", "adaptive"]:
        payload = {**base_payload, "chat_template_kwargs": {"thinking_mode": mode}}
        results.append(check_non_stream(f"thinking_mode={mode}", payload))
        results.append(check_stream(f"thinking_mode={mode}", payload))

    passed = sum(results)
    print(f"\n=== {passed}/{len(results)} endpoint checks passed ===")
    sys.exit(0 if passed == len(results) else 1)


if __name__ == "__main__":
    main()