#!/usr/bin/env python3
"""Written for https://paraloncloud.com/resources/openai-responses-api-open-models
Needs: pip install openai openai-agents pydantic; PRL_KEY=prlc_... in the environment.
A full run makes about 350 requests and costs a few cents at list price.

The OpenAI Responses API against open models on Paralon, measured.

A. SDK feature matrix: which openai-python Responses calls work unchanged.
B. Translation overhead: the same short request through /v1/chat/completions
   and /v1/responses, alternated, latency and TTFT compared.
C. An agent with three tools on 24 support tasks with a known right action,
   run three ways: a hand-written Responses loop, the OpenAI Agents SDK, and a
   Chat Completions loop -- correctness, turns, tokens, cost, time.
"""
import asyncio, json, os, statistics, sys, time, base64, io, struct, zlib
from openai import OpenAI, AsyncOpenAI, BadRequestError
from pydantic import BaseModel
from typing import Literal

BASE = "https://paraloncloud.com/v1"; KEY = os.environ["PRL_KEY"]; MODEL = os.environ.get("MODEL", "qwen3.8-27b")
PRICE_IN, PRICE_OUT = 0.12, 1.70
client = OpenAI(base_url=BASE, api_key=KEY)
aclient = AsyncOpenAI(base_url=BASE, api_key=KEY)
LOW = {"effort": "low"}
res = {}

def red_png(n=64):
    raw = b"".join(b"\x00" + bytes([220, 20, 20]) * n for _ in range(n))
    def chunk(t, d): return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d) & 0xffffffff)
    return b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", n, n, 8, 2, 0, 0, 0)) + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b"")

# ---------------- A. feature matrix ----------------
def feature(name, fn):
    t = time.time()
    try:
        detail = fn()
        res.setdefault("A", []).append({"feature": name, "ok": True, "detail": detail, "s": round(time.time() - t, 2)})
    except Exception as e:
        res.setdefault("A", []).append({"feature": name, "ok": False, "detail": f"{type(e).__name__}: {str(e)[:160]}", "s": round(time.time() - t, 2)})
    print(res["A"][-1], file=sys.stderr)

def f_create():
    r = client.responses.create(model=MODEL, input="Reply with the single word OK.", reasoning=LOW, max_output_tokens=10)
    assert r.output_text.strip(), "empty output_text"
    return f"output_text={r.output_text!r} usage={r.usage.input_tokens}/{r.usage.output_tokens}"

def f_instructions():
    r = client.responses.create(model=MODEL, instructions="Answer in French only.", input="Say good morning.", reasoning=LOW, max_output_tokens=20)
    return r.output_text

def f_vision():
    url = "data:image/png;base64," + base64.b64encode(red_png()).decode()
    r = client.responses.create(model=MODEL, reasoning=LOW, max_output_tokens=10, input=[{"role": "user", "content": [
        {"type": "input_text", "text": "What single colour fills this image? One word."},
        {"type": "input_image", "image_url": url}]}])
    assert "red" in r.output_text.lower(), r.output_text
    return r.output_text

class Ticket(BaseModel):
    intent: Literal["refund", "cancel", "order_status", "other"]
    urgent: bool
    summary: str

def f_parse():
    r = client.responses.parse(model=MODEL, reasoning=LOW, text_format=Ticket,
        input="Classify: 'Cancel order A1002 right now, I ordered the wrong size and it ships tomorrow!'")
    p = r.output_parsed
    assert p is not None and p.intent == "cancel", p
    return p.model_dump()

def f_stream():
    t0 = time.time(); first = None; types = set(); text = ""
    with client.responses.stream(model=MODEL, input="Count from 1 to 5.", reasoning=LOW, max_output_tokens=30) as s:
        for ev in s:
            types.add(ev.type)
            if ev.type == "response.output_text.delta":
                if first is None: first = time.time() - t0
                text += ev.delta
        final = s.get_final_response()
    assert final.output_text == text, (final.output_text, text)
    return f"ttft={first:.2f}s events={sorted(types)}"

TOOLS_WEATHER = [{"type": "function", "name": "get_weather", "description": "Get the weather for a city",
                  "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], "additionalProperties": False}, "strict": True}]

def f_tool_loop():
    r = client.responses.create(model=MODEL, reasoning=LOW, tools=TOOLS_WEATHER, input="What's the weather in Lisbon?")
    calls = [o for o in r.output if o.type == "function_call"]
    assert calls, f"no function_call: {r.output}"
    c = calls[0]
    r2 = client.responses.create(model=MODEL, reasoning=LOW, tools=TOOLS_WEATHER, input=[
        {"role": "user", "content": "What's the weather in Lisbon?"}, c.model_dump(exclude_none=True),
        {"type": "function_call_output", "call_id": c.call_id, "output": json.dumps({"city": "Lisbon", "temp_c": 24, "sky": "sunny"})}])
    assert "24" in r2.output_text or "sunny" in r2.output_text.lower(), r2.output_text
    return f"args={c.arguments} -> {r2.output_text[:80]!r}"

def f_tool_choice():
    r = client.responses.create(model=MODEL, reasoning=LOW, tools=TOOLS_WEATHER, tool_choice={"type": "function", "name": "get_weather"}, input="Tell me a joke.")
    assert any(o.type == "function_call" for o in r.output), r.output
    return "forced call made"

def f_incomplete():
    r = client.responses.create(model=MODEL, reasoning=LOW, input="Write a long essay about rivers.", max_output_tokens=8)
    assert r.status == "incomplete" and r.incomplete_details.reason == "max_output_tokens", (r.status, r.incomplete_details)
    return r.status

def f_prev_id_refused():
    try:
        client.responses.create(model=MODEL, input="hi", previous_response_id="resp_x")
    except BadRequestError as e:
        return "BadRequestError: " + str(e)[:90]
    raise AssertionError("accepted")

def f_reasoning_tokens():
    q = "A bat and a ball cost $1.10 in total. The bat costs $1.00 more than the ball. How much is the ball? Answer with the amount."
    lo = client.responses.create(model=MODEL, input=q, reasoning=LOW, max_output_tokens=2000)
    hi = client.responses.create(model=MODEL, input=q, max_output_tokens=4000)
    res["reasoning"] = {"low_out": lo.usage.output_tokens, "default_out": hi.usage.output_tokens, "low_text": lo.output_text[:60], "default_text": hi.output_text[:60]}
    return res["reasoning"]

# ---------------- B. overhead ----------------
def overhead(n=20):
    lat = {"chat": [], "responses": []}; ttft = {"chat": [], "responses": []}
    msg = "Name the capital of France in one word."
    for i in range(n):
        for kind in ("chat", "responses") if i % 2 == 0 else ("responses", "chat"):
            t = time.time()
            if kind == "chat":
                client.chat.completions.create(model=MODEL, messages=[{"role": "user", "content": msg}], max_tokens=10, temperature=0, extra_body={"chat_template_kwargs": {"enable_thinking": False}})
            else:
                client.responses.create(model=MODEL, input=msg, max_output_tokens=10, temperature=0, reasoning=LOW)
            lat[kind].append(time.time() - t)
        for kind in ("chat", "responses"):
            t = time.time(); first = None
            if kind == "chat":
                for ch in client.chat.completions.create(model=MODEL, messages=[{"role": "user", "content": "Count to 10."}], max_tokens=40, temperature=0, stream=True, extra_body={"chat_template_kwargs": {"enable_thinking": False}}):
                    if first is None and ch.choices and ch.choices[0].delta.content: first = time.time() - t
            else:
                for ev in client.responses.create(model=MODEL, input="Count to 10.", max_output_tokens=40, temperature=0, reasoning=LOW, stream=True):
                    if first is None and ev.type == "response.output_text.delta": first = time.time() - t
            if first: ttft[kind].append(first)
    med = lambda v: round(statistics.median(v), 3)
    res["B"] = {"n": n, "chat_p50_s": med(lat["chat"]), "responses_p50_s": med(lat["responses"]),
                "chat_ttft_p50_s": med(ttft["chat"]), "responses_ttft_p50_s": med(ttft["responses"])}
    print(res["B"], file=sys.stderr)

# ---------------- C. agent ----------------
ORDERS = {
 "A1001": {"status": "shipped", "item": "hiking boots", "total": 129.0, "delivered_days_ago": None},
 "A1002": {"status": "processing", "item": "rain jacket", "total": 89.0, "delivered_days_ago": None},
 "A1003": {"status": "delivered", "item": "tent", "total": 240.0, "delivered_days_ago": 41},
 "A1004": {"status": "delivered", "item": "headlamp", "total": 35.0, "delivered_days_ago": 6},
 "A1005": {"status": "processing", "item": "sleeping bag", "total": 150.0, "delivered_days_ago": None},
 "A1006": {"status": "delivered", "item": "trekking poles", "total": 60.0, "delivered_days_ago": 12},
 "A1007": {"status": "cancelled", "item": "water filter", "total": 45.0, "delivered_days_ago": None},
 "A1008": {"status": "shipped", "item": "backpack", "total": 110.0, "delivered_days_ago": None},
}
POLICY = ("You are a support agent for an outdoor store. Always look the order up with get_order before acting. "
          "Policy: an order can be cancelled only while its status is 'processing'. A refund is allowed only for a delivered order "
          "delivered 30 days ago or less. Never cancel or refund outside the policy; explain why instead. "
          "If the customer gives no order number, ask for it and call no tool. Be brief.")
# (message, expected mutating actions)
TASKS = [
 ("Please cancel order A1002, I changed my mind.", {("cancel_order", "A1002")}),
 ("Cancel A1001 please.", set()),
 ("I want my money back for order A1004, the headlamp is too dim.", {("refund_order", "A1004")}),
 ("Refund order A1003, the tent leaks.", set()),
 ("Where is my order A1008?", set()),
 ("cancel my sleeping bag order A1005", {("cancel_order", "A1005")}),
 ("I'd like a refund for A1006, poles are the wrong length.", {("refund_order", "A1006")}),
 ("Can you refund order A1007?", set()),
 ("Please cancel my order.", set()),
 ("What's the status of A1003?", set()),
 ("Order A1002: don't send it, cancel it.", {("cancel_order", "A1002")}),
 ("I returned the trekking poles from A1006 and want the refund.", {("refund_order", "A1006")}),
 ("A1001 is taking forever, cancel it and refund me.", set()),
 ("Refund A1004 please, and also cancel A1005.", {("refund_order", "A1004"), ("cancel_order", "A1005")}),
 ("Hi, I need a refund.", set()),
 ("Cancel A1008.", set()),
 ("Is A1005 already shipped? If not, cancel it.", {("cancel_order", "A1005")}),
 ("My headlamp from A1004 broke, refund please.", {("refund_order", "A1004")}),
 ("The tent from A1003 arrived damaged a month and a half ago, I want a refund.", set()),
 ("Please cancel A1007.", set()),
 ("Stornează comanda A1002, te rog.", {("cancel_order", "A1002")}),
 ("Bitte erstatten Sie mir Bestellung A1006.", {("refund_order", "A1006")}),
 ("¿Me pueden reembolsar el pedido A1003?", set()),
 ("取消订单 A1005", {("cancel_order", "A1005")}),
]

def run_tool(name, args, actions):
    oid = str(args.get("order_id", "")).upper()
    if name == "get_order":
        o = ORDERS.get(oid)
        return json.dumps({"order_id": oid, **o} if o else {"error": "not found"})
    actions.add((name, oid))
    return json.dumps({"ok": True, "order_id": oid})

def tool_schema(name, desc, extra=None):
    props = {"order_id": {"type": "string"}}
    req = ["order_id"]
    if extra:
        props.update(extra); req += list(extra)
    return {"name": name, "description": desc, "parameters": {"type": "object", "properties": props, "required": req, "additionalProperties": False}}

SCHEMAS = [tool_schema("get_order", "Look up an order by id"),
           tool_schema("cancel_order", "Cancel an order"),
           tool_schema("refund_order", "Refund an order", {"reason": {"type": "string"}})]

async def loop_responses(msg):
    actions, items, turns, tin, tout = set(), [{"role": "user", "content": msg}], 0, 0, 0
    tools = [{"type": "function", **s, "strict": True} for s in SCHEMAS]
    while turns < 8:
        turns += 1
        r = await aclient.responses.create(model=MODEL, instructions=POLICY, input=items, tools=tools, reasoning=LOW, temperature=0, max_output_tokens=600)
        tin += r.usage.input_tokens; tout += r.usage.output_tokens
        calls = [o for o in r.output if o.type == "function_call"]
        if not calls: return actions, turns, tin, tout, r.output_text
        for o in r.output:
            items.append(o.model_dump(exclude_none=True))
        for c in calls:
            items.append({"type": "function_call_output", "call_id": c.call_id, "output": run_tool(c.name, json.loads(c.arguments or "{}"), actions)})
    return actions, turns, tin, tout, "(turn limit)"

async def loop_chat(msg):
    actions, msgs, turns, tin, tout = set(), [{"role": "system", "content": POLICY}, {"role": "user", "content": msg}], 0, 0, 0
    tools = [{"type": "function", "function": s} for s in SCHEMAS]
    while turns < 8:
        turns += 1
        r = await aclient.chat.completions.create(model=MODEL, messages=msgs, tools=tools, temperature=0, max_tokens=600, extra_body={"chat_template_kwargs": {"enable_thinking": False}})
        tin += r.usage.prompt_tokens; tout += r.usage.completion_tokens
        m = r.choices[0].message
        if not m.tool_calls: return actions, turns, tin, tout, m.content or ""
        msgs.append(m.model_dump(exclude_none=True))
        for tc in m.tool_calls:
            msgs.append({"role": "tool", "tool_call_id": tc.id, "content": run_tool(tc.function.name, json.loads(tc.function.arguments or "{}"), actions)})
    return actions, turns, tin, tout, "(turn limit)"

async def loop_agents_sdk(msg):
    from agents import Agent, Runner, function_tool, OpenAIResponsesModel, ModelSettings, set_tracing_disabled
    from openai.types.shared import Reasoning
    set_tracing_disabled(True)
    actions = set()
    @function_tool
    def get_order(order_id: str) -> str:
        """Look up an order by id"""
        return run_tool("get_order", {"order_id": order_id}, actions)
    @function_tool
    def cancel_order(order_id: str) -> str:
        """Cancel an order"""
        return run_tool("cancel_order", {"order_id": order_id}, actions)
    @function_tool
    def refund_order(order_id: str, reason: str) -> str:
        """Refund an order"""
        return run_tool("refund_order", {"order_id": order_id}, actions)
    agent = Agent(name="support", instructions=POLICY, tools=[get_order, cancel_order, refund_order],
                  model=OpenAIResponsesModel(model=MODEL, openai_client=aclient),
                  model_settings=ModelSettings(temperature=0, max_tokens=600, reasoning=Reasoning(effort="low")))
    r = await Runner.run(agent, msg, max_turns=8)
    tin = sum(u.usage.input_tokens for u in r.raw_responses if u.usage) if r.raw_responses else 0
    tout = sum(u.usage.output_tokens for u in r.raw_responses if u.usage) if r.raw_responses else 0
    return actions, len(r.raw_responses), tin, tout, str(r.final_output)

async def agent_bench(name, fn, par=6):
    sem = asyncio.Semaphore(par); rows = []
    async def one(i, msg, want):
        async with sem:
            t = time.time()
            try:
                got, turns, tin, tout, text = await fn(msg)
                rows.append({"i": i, "ok": got == want, "got": sorted(got), "want": sorted(want), "turns": turns, "in": tin, "out": tout, "s": time.time() - t, "text": text[:100]})
            except Exception as e:
                rows.append({"i": i, "ok": False, "error": f"{type(e).__name__}: {str(e)[:200]}", "turns": 0, "in": 0, "out": 0, "s": time.time() - t})
    t0 = time.time()
    await asyncio.gather(*[one(i, m, w) for i, (m, w) in enumerate(TASKS)])
    wall = time.time() - t0
    ok = sum(r["ok"] for r in rows); tin = sum(r["in"] for r in rows); tout = sum(r["out"] for r in rows)
    summary = {"variant": name, "tasks": len(rows), "correct": ok, "errors": sum(1 for r in rows if "error" in r),
               "turns_avg": round(sum(r["turns"] for r in rows) / len(rows), 2), "in_per_task": round(tin / len(rows)), "out_per_task": round(tout / len(rows)),
               "s_per_task_p50": round(statistics.median(r["s"] for r in rows), 2), "wall_s": round(wall, 1),
               "usd_per_1k_tasks": round((tin * PRICE_IN + tout * PRICE_OUT) / 1e6 / len(rows) * 1000, 3),
               "wrong": [r for r in sorted(rows, key=lambda r: r["i"]) if not r["ok"]]}
    res.setdefault("C", []).append(summary)
    print(json.dumps({k: v for k, v in summary.items() if k != "wrong"}), file=sys.stderr)

def main():
    for name, fn in [("responses.create + output_text", f_create), ("instructions", f_instructions), ("input_image (vision)", f_vision),
                     ("responses.parse with a Pydantic model", f_parse), ("responses.stream + get_final_response", f_stream),
                     ("function tool -> function_call_output -> answer", f_tool_loop), ("tool_choice forcing a function", f_tool_choice),
                     ("max_output_tokens -> status incomplete", f_incomplete), ("previous_response_id refused (BadRequestError)", f_prev_id_refused),
                     ("reasoning.effort low vs default", f_reasoning_tokens)]:
        feature(name, fn)
    overhead()
    for name, fn in [("responses loop", loop_responses), ("agents sdk", loop_agents_sdk), ("chat completions loop", loop_chat)]:
        asyncio.run(agent_bench(name, fn))
    json.dump(res, open("results.json", "w"), indent=1, ensure_ascii=False, default=str)
    print(json.dumps(res, indent=1, ensure_ascii=False, default=str))

main()