"""Raeto routing for any OpenAI-compatible Python client. Standard library only; Python 3.8+.

    from openai import OpenAI
    from raeto_route import with_raeto
    client = with_raeto(OpenAI(), api_key="sk-raeto-...")
    # use client exactly as before

Before each chat request, Raeto picks the cheapest model that can handle it; your client then calls
that model with your own provider key, as usual. Raeto never sees your provider key or the reply.
If Raeto is slow or unreachable, the request goes to the model you asked for, so adding Raeto can
never break your app. Works with OpenAI and AsyncOpenAI.

Tiers: by default `top` and `mid` are the model you asked for and `fast` is a cheap model from the
same family. Pass `models={"fast": ..., "mid": ..., "top": ...}` to choose your own.
"""
import asyncio
import inspect
import json
import re
import urllib.request

__all__ = ["route", "with_raeto"]
__version__ = "1.0.0"
DEFAULT_BASE = "https://api.raeto.cc"


def _default_tiers(model):
    m = str(model or "")
    fast = None
    if re.match(r"^(gpt-|o\d|chatgpt-)", m, re.I) or m.lower().startswith("openai/"):
        fast = "openai/gpt-4o-mini" if m.lower().startswith("openai/") else "gpt-4o-mini"
    elif "claude" in m.lower():
        fast = "anthropic/claude-3.5-haiku" if m.lower().startswith("anthropic/") else "claude-3-5-haiku-latest"
    elif "gemini" in m.lower():
        fast = "google/gemini-2.5-flash" if m.lower().startswith("google/") else "gemini-2.5-flash"
    return {"fast": fast or m, "mid": m, "top": m}


def route(request, api_key, models=None, floor=None, base=DEFAULT_BASE, timeout=1.5):
    """Ask Raeto which model to use. Returns Raeto's answer (a dict with "model", "tier", ...) or None."""
    if not api_key or not isinstance(request, dict) or not isinstance(request.get("messages"), list):
        return None
    body = {"messages": request["messages"]}
    for k in ("tools", "max_tokens"):
        if request.get(k):
            body[k] = request[k]
    if models:
        body["models"] = models
    if floor:
        body["floor"] = floor
    req = urllib.request.Request(
        base.rstrip("/") + "/v1/route", data=json.dumps(body).encode(), method="POST",
        headers={"Authorization": "Bearer " + api_key, "Content-Type": "application/json", "User-Agent": "raeto-route-python/" + __version__},
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            j = json.loads(r.read().decode())
        return j if isinstance(j, dict) and isinstance(j.get("model"), str) else None
    except Exception:
        return None


def with_raeto(client, api_key, models=None, floor=None, base=DEFAULT_BASE, timeout=1.5, on_route=None):
    """Wrap an OpenAI-compatible client so every chat.completions.create call is routed first.
    Skip routing for one call with raeto=False."""
    completions = client.chat.completions
    if getattr(completions, "_raeto", False):
        return client
    original = completions.create

    def _pick(params):
        rest = {k: v for k, v in params.items() if k != "raeto"}
        if params.get("raeto") is False:
            return rest
        decision = route(rest, api_key, models or _default_tiers(rest.get("model")), floor, base, timeout)
        if on_route:
            try:
                on_route(decision, rest)
            except Exception:
                pass
        if decision:
            rest["model"] = decision["model"]
        return rest

    if inspect.iscoroutinefunction(original):
        async def create(*args, **params):
            to_thread = getattr(asyncio, "to_thread", None)
            rest = await to_thread(_pick, params) if to_thread else _pick(params)
            return await original(*args, **rest)
    else:
        def create(*args, **params):
            return original(*args, **_pick(params))

    completions.create = create
    completions._raeto = True
    return client