"""OpenAI-compatible chat-completions client used by every instrument.

One thin wrapper around urllib + json. No deps. Every instrument imports
`call_model` from here so endpoint/auth/model config lives in exactly one
place (loaded from .env at import time, overridable via env vars).

Env vars (all optional — sensible defaults):
  ENDPOINT   chat-completions URL       (default http://127.0.0.1:1234/v1/chat/completions)
  MODEL      model id                    (default "local")
  API_KEY    bearer token, sent if set   (default none)
  MAX_TOKENS per-turn generation cap      (default 8000)
  TEMPERATURE sampling temp              (default 0.0)
  HTTP_TIMEOUT per-request wall timeout   (default 900s)
  NOTHINK    "1" = append empty <think/>   (default off)
"""
import json
import os
import time
import urllib.request
import urllib.error

ENDPOINT = os.environ.get(
    "ENDPOINT", "http://127.0.0.1:1234/v1/chat/completions")
MODEL = os.environ.get("MODEL", "local")
API_KEY = os.environ.get("API_KEY", "")
MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "8000"))
TEMPERATURE = float(os.environ.get("TEMPERATURE", "0"))
HTTP_TIMEOUT = int(os.environ.get("HTTP_TIMEOUT", "900"))
NOTHINK = os.environ.get("NOTHINK", "") == "1"


def call_model(messages, tools=None, *, max_tokens=None, temperature=None,
               model=None, endpoint=None, api_key=None, http_timeout=None):
    """Send a chat-completions request. Returns (resp_dict, elapsed_seconds).

    All kwargs default to the module-level config (env-driven). Pass kwargs
    only when an instrument needs to override per-call (e.g. router_bench
    wants NOTHINK but coding-suite does not).
    """
    msgs = list(messages)
    if NOTHINK:
        msgs = msgs + [{"role": "assistant", "content": "<think></think>"}]
    body = {
        "model": model or MODEL,
        "messages": msgs,
        "max_tokens": max_tokens if max_tokens is not None else MAX_TOKENS,
        "temperature": temperature if temperature is not None else TEMPERATURE,
    }
    if tools:
        body["tools"] = tools
    data = json.dumps(body).encode()
    headers = {"Content-Type": "application/json"}
    key = api_key if api_key is not None else API_KEY
    if key:
        headers["Authorization"] = "Bearer " + key
    req = urllib.request.Request(
        endpoint or ENDPOINT, data=data, headers=headers)
    t0 = time.time()
    timeout = http_timeout or HTTP_TIMEOUT
    with urllib.request.urlopen(req, timeout=timeout) as r:
        resp = json.load(r)
    return resp, time.time() - t0


def content(resp):
    """Extract the assistant content string from a chat-completions response."""
    return resp["choices"][0]["message"].get("content") or ""


def tool_calls(resp):
    """Extract the tool_calls list (empty if none)."""
    return resp["choices"][0]["message"].get("tool_calls") or []


def reasoning(resp):
    """Extract reasoning_content (thinking) if the backend exposes it."""
    msg = resp["choices"][0]["message"]
    return msg.get("reasoning_content") or ""


def finish_reason(resp):
    return resp["choices"][0].get("finish_reason")


def usage(resp):
    """Return a normalized usage dict: prompt/completion/reasoning tokens."""
    u = resp.get("usage") or {}
    return {
        "prompt_tokens": u.get("prompt_tokens") or 0,
        "completion_tokens": u.get("completion_tokens") or 0,
        "reasoning_tokens":
            (u.get("completion_tokens_details") or {}).get("reasoning_tokens") or 0,
    }


def timings(resp):
    """Return prefill/decode tok/s if the backend reports them (llama-server).

    OpenAI-style responses don't have these; they'll be None.
    """
    t = resp.get("timings") or {}
    return {
        "prefill_ps": t.get("prompt_per_second"),
        "decode_ps": t.get("predicted_per_second"),
    }