cp -r supportdesk ~/work/m13 && cd ~/work/m13
python3.11 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
pip install scikit-learn==1.9.1 scipy==1.17.1 matplotlib==3.11.2
export PYTHONPATH=.
python -m pytest -q tests/test_m13_gateway.py tests/test_m13_ops.py

"""A small LLM gateway: one front door in front of several chat providers.

Every provider is a callable with the same signature as supportdesk.llm.chat,
so a real provider (llm.chat with provider="groq") and a fake one used in
tests are interchangeable. The gateway adds what production needs and a single
call does not have: fallback chains, retries with exponential backoff and
jitter that honor Retry-After, per-user quotas, a kill switch, an exact
response cache, and prefix-cache statistics.

Time is injectable (clock and sleep), so the demo and tests run instantly and
print the same numbers every time.
"""
from __future__ import annotations

import hashlib
import json
import random
import time
from collections import OrderedDict, defaultdict
from dataclasses import dataclass, field
from typing import Any, Callable

from supportdesk.llm import ChatResult
from supportdesk.pricing import PRICES, cost_usd
from supportdesk.stand_in import ScriptedLLM
from supportdesk.tokens import count_messages

ChatFn = Callable[..., ChatResult]


# Errors ------------------------------------------------------------------------

class ProviderError(Exception):
    """A failed provider call, already classified as retryable or not."""

    retryable = False

    def __init__(self, message: str, status: int | None = None, retry_after: float | None = None) -> None:
        super().__init__(message)
        self.status = status
        self.retry_after = retry_after


class RateLimited(ProviderError):
    retryable = True


class ServerError(ProviderError):
    retryable = True


class BadRequest(ProviderError):
    retryable = False


class QuotaExceeded(Exception):
    pass


class KillSwitchOn(Exception):
    pass


class AllProvidersFailed(Exception):
    pass


def classify(exc: Exception) -> ProviderError:
    """Map an exception from the openai SDK (or a fake) onto the gateway's error types."""
    if isinstance(exc, ProviderError):
        return exc
    import openai  # imported here so fakes never need the SDK

    if isinstance(exc, openai.RateLimitError):
        header = exc.response.headers.get("retry-after") if exc.response is not None else None
        try:
            retry_after = float(header) if header is not None else None
        except ValueError:
            retry_after = None  # an HTTP date instead of seconds; fall back to our own backoff
        return RateLimited(str(exc), 429, retry_after)
    if isinstance(exc, (openai.APITimeoutError, openai.APIConnectionError)):
        return ServerError(str(exc), None)
    if isinstance(exc, openai.APIStatusError):
        if exc.status_code >= 500:
            return ServerError(str(exc), exc.status_code)
        return BadRequest(str(exc), exc.status_code)
    raise exc


# Configuration -------------------------------------------------------------------

@dataclass(frozen=True)
class Route:
    """One step of a fallback chain: a provider name, the callable, and the model it serves."""

    provider: str
    call: ChatFn
    model: str


@dataclass(frozen=True)
class RetryPolicy:
    max_attempts: int = 3          # tries per provider before falling back
    base_s: float = 0.5            # first backoff ceiling
    cap_s: float = 8.0             # backoff never exceeds this
    max_retry_after_s: float = 10.0  # a longer Retry-After means "go to the next provider"
    deadline_s: float = 30.0       # total time budget for one gateway call

    def backoff(self, attempt: int, rng: random.Random) -> float:
        """Full jitter: a random wait between 0 and min(cap, base * 2**attempt)."""
        return rng.uniform(0, min(self.cap_s, self.base_s * 2 ** attempt))


@dataclass
class Quota:
    usd_per_day: float = 0.50
    requests_per_minute: int = 10


class QuotaBook:
    """Per-user spend per day and requests per minute, checked before and charged after a call."""

    def __init__(self, default: Quota, clock: Callable[[], float]) -> None:
        self.default = default
        self.overrides: dict[str, Quota] = {}
        self.clock = clock
        self.spent: dict[tuple[str, int], float] = defaultdict(float)
        self.recent: dict[str, list[float]] = defaultdict(list)

    def _day(self) -> int:
        return int(self.clock() // 86_400)

    def check(self, user: str) -> None:
        quota = self.overrides.get(user, self.default)
        now = self.clock()
        self.recent[user] = [t for t in self.recent[user] if now - t < 60]
        if len(self.recent[user]) >= quota.requests_per_minute:
            raise QuotaExceeded(f"{user}: {quota.requests_per_minute} requests per minute reached")
        if self.spent[(user, self._day())] >= quota.usd_per_day:
            raise QuotaExceeded(f"{user}: daily budget {quota.usd_per_day:.2f} USD reached")
        self.recent[user].append(now)

    def charge(self, user: str, usd: float) -> None:
        self.spent[(user, self._day())] += usd


class ExactCache:
    """Response cache keyed on the full request. Least recently used entries are evicted first."""

    def __init__(self, max_entries: int = 1000, ttl_s: float = 3600, clock: Callable[[], float] = time.monotonic) -> None:
        self.max_entries, self.ttl_s, self.clock = max_entries, ttl_s, clock
        self.store: OrderedDict[str, tuple[float, ChatResult]] = OrderedDict()
        self.hits = self.misses = 0

    @staticmethod
    def key(route: str, messages: list[dict[str, Any]], kwargs: dict[str, Any]) -> str:
        payload = json.dumps({"route": route, "messages": messages, "kwargs": kwargs}, sort_keys=True, default=str)
        return hashlib.sha256(payload.encode()).hexdigest()

    def get(self, key: str) -> ChatResult | None:
        item = self.store.get(key)
        if item is None or self.clock() - item[0] > self.ttl_s:
            self.store.pop(key, None)
            self.misses += 1
            return None
        self.store.move_to_end(key)
        self.hits += 1
        return item[1]

    def put(self, key: str, result: ChatResult) -> None:
        self.store[key] = (self.clock(), result)
        self.store.move_to_end(key)
        while len(self.store) > self.max_entries:
            self.store.popitem(last=False)


class PrefixStats:
    """Tracks whether a request's stable prefix (every message but the last) was seen recently.

    Providers with prompt caching bill a recently seen prefix at the cached rate.
    The gateway cannot see the provider's cache, but it can measure how often
    your traffic *could* hit it, which is what prompt layout decides.
    """

    def __init__(self, ttl_s: float = 300, clock: Callable[[], float] = time.monotonic) -> None:
        self.ttl_s, self.clock = ttl_s, clock
        self.seen: dict[str, float] = {}
        self.requests = self.hits = self.prefix_tokens_hit = self.input_tokens = 0

    def observe(self, messages: list[dict[str, Any]]) -> bool:
        prefix = messages[:-1]
        key = hashlib.sha256(json.dumps(prefix, sort_keys=True).encode()).hexdigest()
        now = self.clock()
        hit = bool(prefix) and key in self.seen and now - self.seen[key] <= self.ttl_s
        self.seen[key] = now
        self.requests += 1
        self.input_tokens += count_messages(messages)
        if hit:
            self.hits += 1
            self.prefix_tokens_hit += count_messages(prefix) - 3  # minus the reply priming tokens
        return hit


@dataclass
class GatewayResult:
    result: ChatResult
    route: str
    provider: str
    cache: str                     # "exact" or "miss"
    prefix_hit: bool
    cost_usd: float
    attempts: list[dict[str, Any]] = field(default_factory=list)


# The gateway ------------------------------------------------------------------------

class Gateway:
    def __init__(self, routes: dict[str, list[Route]], *, policy: RetryPolicy | None = None,
                 quota: Quota | None = None, clock: Callable[[], float] = time.monotonic,
                 sleep: Callable[[float], None] = time.sleep, seed: int = 0,
                 on_event: Callable[[dict[str, Any]], None] | None = None) -> None:
        self.routes = routes
        self.policy = policy or RetryPolicy()
        self.clock, self.sleep = clock, sleep
        self.rng = random.Random(seed)
        self.quotas = QuotaBook(quota or Quota(), clock)
        self.cache = ExactCache(clock=clock)
        self.prefix = PrefixStats(clock=clock)
        self.killed: dict[str, str] = {}          # route name or "*" -> reason
        self.disabled_providers: set[str] = set()
        self.on_event = on_event or (lambda event: None)

    # Operator controls
    def kill(self, route: str = "*", reason: str = "") -> None:
        self.killed[route] = reason or "no reason given"

    def revive(self, route: str = "*") -> None:
        self.killed.pop(route, None)

    def disable_provider(self, provider: str) -> None:
        self.disabled_providers.add(provider)

    def enable_provider(self, provider: str) -> None:
        self.disabled_providers.discard(provider)

    # The one method callers use
    def chat(self, messages: list[dict[str, Any]], *, route: str, user: str, **kwargs: Any) -> GatewayResult:
        for name in ("*", route):
            if name in self.killed:
                self.on_event({"type": "killed", "route": route, "reason": self.killed[name]})
                raise KillSwitchOn(f"route {route!r} is switched off: {self.killed[name]}")
        self.quotas.check(user)
        prefix_hit = self.prefix.observe(messages)

        cacheable = kwargs.get("temperature", 0.0) in (0, 0.0) and not kwargs.get("tools")
        key = ExactCache.key(route, messages, kwargs) if cacheable else ""
        if cacheable:
            cached = self.cache.get(key)
            if cached is not None:
                self.on_event({"type": "cache_hit", "route": route})
                return GatewayResult(cached, route, cached.provider, "exact", prefix_hit, 0.0)

        started = self.clock()
        attempts: list[dict[str, Any]] = []
        for step in self.routes[route]:
            if step.provider in self.disabled_providers:
                attempts.append({"provider": step.provider, "outcome": "disabled"})
                continue
            for attempt in range(self.policy.max_attempts):
                if self.clock() - started > self.policy.deadline_s:
                    raise AllProvidersFailed(f"deadline {self.policy.deadline_s}s exceeded; attempts={attempts}")
                try:
                    result = step.call(messages, model=step.model, **kwargs)
                except Exception as exc:  # noqa: BLE001  (classify re-raises what it does not know)
                    error = classify(exc)
                    record = {"provider": step.provider, "attempt": attempt + 1,
                              "outcome": type(error).__name__, "status": error.status}
                    attempts.append(record)
                    self.on_event({"type": "error", "route": route, **record})
                    if not error.retryable:
                        raise
                    if error.retry_after is not None and error.retry_after > self.policy.max_retry_after_s:
                        record["action"] = f"retry-after {error.retry_after:g}s too long, fall back"
                        break
                    if attempt + 1 == self.policy.max_attempts:
                        record["action"] = "attempts used up, fall back"
                        break
                    delay = self.policy.backoff(attempt, self.rng)
                    if error.retry_after is not None:
                        delay = max(delay, error.retry_after)
                    if self.clock() - started + delay > self.policy.deadline_s:
                        record["action"] = "wait would break deadline, fall back"
                        break
                    record["action"] = f"sleep {delay:.2f}s"
                    self.sleep(delay)
                    continue
                usd = cost_usd(result.usage, step.model) if step.model in PRICES else 0.0
                self.quotas.charge(user, usd)
                attempts.append({"provider": step.provider, "attempt": attempt + 1, "outcome": "ok"})
                self.on_event({"type": "ok", "route": route, "provider": step.provider,
                               "latency_ms": result.latency_ms, "usd": usd})
                if cacheable:
                    self.cache.put(key, result)
                return GatewayResult(result, route, step.provider, "miss", prefix_hit, usd, attempts)
        raise AllProvidersFailed(f"every provider failed for route {route!r}: {attempts}")


# Fakes for demos and tests ----------------------------------------------------------

class FakeClock:
    """Simulated time: sleep() advances the clock instantly and remembers each wait."""

    def __init__(self, start: float = 1_000_000.0) -> None:
        self.now = start
        self.sleeps: list[float] = []

    def __call__(self) -> float:
        return self.now

    def sleep(self, seconds: float) -> None:
        self.sleeps.append(round(seconds, 3))
        self.now += seconds


class FakeProvider:
    """A provider that fails on a script, then answers through ScriptedLLM (not a model).

    script entries: "ok", "429" (no Retry-After), "429:2.5" (Retry-After 2.5 s),
    "500", "timeout", "400". After the script runs out, outcome_fn decides (or "ok").
    latency_fn, if given, draws each call's latency in ms (for realistic logs).
    """

    def __init__(self, name: str, script: list[str] | None = None, latency_ms: float = 400.0,
                 clock: FakeClock | None = None, latency_fn: Callable[[], float] | None = None,
                 outcome_fn: Callable[[], str] | None = None) -> None:
        self.name = name
        self.script = list(script or [])
        self.latency_ms = latency_ms
        self.latency_fn = latency_fn
        self.outcome_fn = outcome_fn
        self.clock = clock
        self.calls = 0
        self.llm = ScriptedLLM(responder=self._respond, model=f"{name}-stand-in")

    def _respond(self, messages: list[dict[str, Any]], kwargs: dict[str, Any]) -> str:
        subject = messages[-1]["content"].splitlines()[0][:60]
        return f"[{self.name}] draft for: {subject}"

    def __call__(self, messages: list[dict[str, Any]], **kwargs: Any) -> ChatResult:
        self.calls += 1
        latency = self.latency_fn() if self.latency_fn else self.latency_ms
        if self.clock is not None:
            self.clock.now += latency / 1000
        outcome = self.script.pop(0) if self.script else (self.outcome_fn() if self.outcome_fn else "ok")
        if outcome.startswith("429"):
            retry_after = float(outcome.split(":")[1]) if ":" in outcome else None
            raise RateLimited(f"{self.name}: 429 Too Many Requests", 429, retry_after)
        if outcome == "500":
            raise ServerError(f"{self.name}: 500 Internal Server Error", 500)
        if outcome == "timeout":
            raise ServerError(f"{self.name}: request timed out", None)
        if outcome == "400":
            raise BadRequest(f"{self.name}: 400 unsupported parameter", 400)
        result = self.llm(messages, **kwargs)
        result.model = kwargs.get("model", result.model)
        result.provider = self.name
        result.latency_ms = round(latency, 1)
        return result


def ticket_messages(subject: str, body: str) -> list[dict[str, str]]:
    """Stable system prompt first, the volatile ticket last (cache-aware layout from Module 7)."""
    system = ("You are Brightlane's support assistant. Draft a short, polite reply for a human agent "
              "to review. Cite help-center articles by id. Never promise refunds; agents decide those.")
    return [{"role": "system", "content": system},
            {"role": "user", "content": f"Subject: {subject}\n\n{body}"}]


def demo() -> None:
    clock = FakeClock()
    groq = FakeProvider("groq", ["429:1.5", "500", "ok"], latency_ms=350, clock=clock)
    gemini = FakeProvider("gemini", [], latency_ms=900, clock=clock)
    local = FakeProvider("ollama", [], latency_ms=2500, clock=clock)
    routes = {"draft": [Route("groq", groq, "openai/gpt-oss-120b"),
                        Route("gemini", gemini, "gemini-3.5-flash"),
                        Route("ollama", local, "qwen3:8b")]}
    gw = Gateway(routes, quota=Quota(usd_per_day=0.002, requests_per_minute=5), clock=clock, sleep=clock.sleep)
    msgs = ticket_messages("Charged twice this month", "My card was charged 288 USD twice. Please refund.")

    print("1) groq fails with 429 (Retry-After 1.5 s), then 500, then succeeds")
    r = gw.chat(msgs, route="draft", user="u-17")
    for a in r.attempts:
        print("  ", a)
    print(f"   served by {r.provider}: {r.result.text!r}  cost={r.cost_usd:.6f} USD")
    print(f"   simulated sleeps={clock.sleeps}  simulated elapsed={clock.now - 1_000_000:.2f} s")

    print("2) the identical request again (temperature 0)")
    r = gw.chat(msgs, route="draft", user="u-17")
    print(f"   cache={r.cache} provider={r.provider} cost={r.cost_usd}  groq calls so far={groq.calls}")

    print("3) groq rate-limited for a long time (Retry-After 30 s): fall back at once")
    groq.script = ["429:30"]
    other = ticket_messages("Export to CSV", "How do I export my tasks to CSV?")
    r = gw.chat(other, route="draft", user="u-17")
    for a in r.attempts:
        print("  ", a)
    print(f"   served by {r.provider}: {r.result.text!r}")

    print("4) one user sends 7 requests in the same minute (limit 5 per minute)")
    for i in range(7):
        try:
            gw.chat(ticket_messages(f"Question {i}", "Where is the billing page?"), route="draft", user="u-99")
            print(f"   request {i}: ok")
        except QuotaExceeded as exc:
            print(f"   request {i}: refused ({exc})")

    print("5) incident: the operator flips the kill switch on the draft route")
    gw.kill("draft", "drafts quoting wrong refund window, INC-311")
    try:
        gw.chat(msgs, route="draft", user="u-17")
    except KillSwitchOn as exc:
        print(f"   refused: {exc}; the ticket goes to the human queue")
    gw.revive("draft")

    print(f"exact cache: hits={gw.cache.hits} misses={gw.cache.misses}; "
          f"prefix seen before on {gw.prefix.hits} of {gw.prefix.requests} requests")


if __name__ == "__main__":
    demo()

"""Wire the gateway to real providers through supportdesk.llm.chat.

Only providers whose API key is set join the chain; Ollama joins when
OLLAMA_BASE_URL is set or USE_OLLAMA=1. Run with no keys and it tells you so.
"""
from __future__ import annotations

import functools
import os
import time

from m13_gateway import AllProvidersFailed, Gateway, Route, ticket_messages
from supportdesk import llm


def build_routes() -> dict[str, list[Route]]:
    chain = []
    for provider in ("groq", "gemini"):
        if os.environ.get(llm.PROVIDERS[provider]["key_env"]):
            chain.append(Route(provider, functools.partial(llm.chat, provider=provider),
                               llm.PROVIDERS[provider]["default_model"]))
    if os.environ.get("OLLAMA_BASE_URL") or os.environ.get("USE_OLLAMA") == "1":
        chain.append(Route("ollama", functools.partial(llm.chat, provider="ollama"), "qwen3:8b"))
    return {"draft": chain}


def main() -> None:
    routes = build_routes()
    print("fallback chain:", [r.provider for r in routes["draft"]] or "empty (set GROQ_API_KEY, GEMINI_API_KEY or USE_OLLAMA=1)")
    if not routes["draft"]:
        return
    gateway = Gateway(routes)
    started = time.perf_counter()
    try:
        r = gateway.chat(ticket_messages("Charged twice this month", "My card was charged 288 USD twice."),
                         route="draft", user="maya", max_tokens=300)
        print(f"served by {r.provider} ({r.result.model}) in {r.result.latency_ms:.0f} ms, cost {r.cost_usd:.6f} USD")
        print(r.result.text)
    except AllProvidersFailed as exc:
        print(f"all providers failed after {time.perf_counter() - started:.1f} s")
        print(str(exc)[:230] + " ...")


if __name__ == "__main__":
    main()