#!/usr/bin/env python3
"""Compare DS4 and Rust endpoints configured for the same decoding mode."""

import argparse
import json
import urllib.request


CASES = [
    "hi",
    "Reply with exactly three words describing a calm sea.",
    "Write the first eight positive odd numbers separated by commas.",
]


def post(base_url, path, payload):
    request = urllib.request.Request(
        base_url.rstrip("/") + path,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(request, timeout=600) as response:
        return json.load(response)


def chat(base_url, prompt):
    result = post(
        base_url,
        "/v1/chat/completions",
        {
            "model": "deepseek-v4-flash",
            "messages": [{"role": "user", "content": prompt}],
            "reasoning_effort": "none",
            "temperature": 0,
            "max_tokens": 64,
        },
    )
    choice = result["choices"][0]
    message = choice["message"]
    return {
        "content": message.get("content"),
        "reasoning": message.get("reasoning_content"),
        "finish_reason": choice.get("finish_reason"),
        "usage": result.get("usage"),
    }


def anthropic(base_url, prompt):
    result = post(
        base_url,
        "/v1/messages",
        {
            "model": "deepseek-v4-flash",
            "messages": [{"role": "user", "content": prompt}],
            "thinking": {"type": "disabled"},
            "temperature": 0,
            "max_tokens": 64,
        },
    )
    return {
        "content": result.get("content"),
        "stop_reason": result.get("stop_reason"),
        "usage": result.get("usage"),
    }


def responses(base_url, prompt):
    result = post(
        base_url,
        "/v1/responses",
        {
            "model": "deepseek-v4-flash",
            "input": prompt,
            "reasoning": {"effort": "none"},
            "temperature": 0,
            "max_output_tokens": 64,
        },
    )
    return {
        "output": normalize(result.get("output")),
        "status": result.get("status"),
        "usage": result.get("usage"),
    }


def normalize(value):
    if isinstance(value, list):
        return [normalize(item) for item in value]
    if isinstance(value, dict):
        return {
            key: normalize(item)
            for key, item in value.items()
            if key not in {"id", "created_at"}
        }
    return value


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("reference_url", help="DS4 reference server URL")
    parser.add_argument("rust_url", help="Rust DS4Server URL in the same mode")
    args = parser.parse_args()

    failures = []
    for prompt in CASES:
        for name, request in (
            ("chat", chat),
            ("anthropic", anthropic),
            ("responses", responses),
        ):
            reference = request(args.reference_url, prompt)
            rust = request(args.rust_url, prompt)
            if reference != rust:
                failures.append(
                    {
                        "case": name,
                        "prompt": prompt,
                        "reference": reference,
                        "rust": rust,
                    }
                )
            else:
                print(f"ok {name}: {prompt}")
    if failures:
        print(json.dumps(failures, indent=2, ensure_ascii=False))
        raise SystemExit(1)
    print("all deterministic DS4/Rust endpoint outputs and usage records match")


if __name__ == "__main__":
    main()