#!/usr/bin/env python3
"""Routing policy for the frank-medium proxy: which replica serves a request.

Kept apart from the HTTP plumbing in affinity-proxy.py so the decisions can be
tested without aiohttp, a network or a GPU — the first version of these rules
shipped with a bug that only appeared on a continuation request, and a test
that loads this module catches that class of mistake directly.

Two cards, two behaviours, and the workload decides which is right:

* A **single-turn** request is one member of a batch (a benchmark run, a burst
  of parallel subagents). Spreading those over both replicas lets both cards
  prefill at once; measured on this host, a 16-request burst of 15k-token
  prompts finishes in 18.6 s spread versus 30.1 s pinned.
* A **continuation** re-sends a prefix the replica has already prefilled. Pinned
  to that replica the turn costs a cache lookup; spread, half the turns pay a
  cold prefill again (measured: 1096 ms versus 2676 ms mean warm-turn TTFT on
  90k-token contexts).

ROUTING picks the strategy; the default `adaptive` applies both rules above.
`affinity` and `roundrobin` exist so the choice can be re-measured, and are not
production settings.
"""

from __future__ import annotations

import collections
import hashlib
import itertools
import json
import os

UPSTREAMS = [u.strip() for u in os.environ.get(
    "UPSTREAMS", "http://127.0.0.1:7771,http://127.0.0.1:7772").split(",") if u.strip()]
ROUTING = os.environ.get("ROUTING", "adaptive").strip().lower()
# Leading bytes of prompt text that identify a conversation's shared prefix.
# 4 KiB covers a system prompt plus the start of the first user turn, which is
# where agent loops keep the text they resend every turn.
AFFINITY_BYTES = int(os.environ.get("AFFINITY_BYTES", "4096"))
# Prefixes remembered for pinning. Bounded so a long-lived proxy cannot grow
# without limit; 4096 conversations is far beyond what this fleet runs at once.
KNOWN_MAX = int(os.environ.get("KNOWN_MAX", "4096"))

_counter = itertools.count()
# digest -> replica index, insertion-ordered for O(1) LRU eviction. A prefix is
# remembered as soon as any request carries it; only a continuation is pinned
# there, because only a continuation can profit from the cache entry.
_known: "collections.OrderedDict[bytes, int]" = collections.OrderedDict()


def affinity_key(body: bytes, remote: str | None) -> tuple[bytes, bool]:
    """Key a request by its shared prefix, and report whether it continues one.

    Chat and completions bodies both carry their prompt text in JSON, so the
    leading bytes of that text are the prefix; a body we cannot parse (or an
    endpoint with no prompt, e.g. /v1/models) falls back to the client address,
    which still beats nothing for a single-client loop.
    """
    fallback = (remote or "anonymous").encode()
    try:
        payload = json.loads(body)
    except Exception:
        return fallback, False

    text = ""
    continuation = False
    if isinstance(payload, dict):
        if isinstance(payload.get("prompt"), str):
            text = payload["prompt"]
        elif isinstance(payload.get("messages"), list):
            continuation = len(payload["messages"]) > 1
            parts = []
            for msg in payload["messages"]:
                if isinstance(msg, dict):
                    content = msg.get("content")
                    if isinstance(content, str):
                        parts.append(content)
                    elif isinstance(content, list):  # multimodal content blocks
                        parts.extend(
                            blk.get("text", "") for blk in content
                            if isinstance(blk, dict) and isinstance(blk.get("text"), str)
                        )
            text = "\n".join(parts)
    if not text:
        return fallback, continuation
    return text[:AFFINITY_BYTES].encode("utf-8", "replace"), continuation


def pick(key: bytes, continuation: bool, up: list[bool]) -> int:
    """Replica for a request, or -1 when none is available.

    Down replicas are never selected — a rollout must not send traffic to an
    engine that is still compiling its graphs — and a pin whose replica is down
    is re-established on a live one rather than kept pointing at a corpse.
    """
    digest = hashlib.sha1(key).digest()

    if ROUTING == "roundrobin":
        start = next(_counter) % len(UPSTREAMS)
    elif ROUTING == "adaptive":
        entry = _known.get(digest)
        if continuation and entry is not None and up[entry]:
            _known.move_to_end(digest)
            return entry
        # Anything that is not a continuation spreads by counter — including a
        # batch of requests that happen to share one prefix. Hashing the prefix
        # would send every member of such a batch to whichever card the hash
        # names (measured: 29.8 s for a 16-request burst, against 18.6 s when
        # spread). The counter is only advanced by these requests, so batch
        # members alternate across the cards.
        start = next(_counter) % len(UPSTREAMS)
    else:  # affinity
        start = int.from_bytes(digest[:8], "big") % len(UPSTREAMS)

    for offset in range(len(UPSTREAMS)):
        idx = (start + offset) % len(UPSTREAMS)
        if up[idx]:
            if ROUTING == "adaptive":
                _known[digest] = idx
                _known.move_to_end(digest)
                if len(_known) > KNOWN_MAX:
                    _known.popitem(last=False)
            return idx
    return -1


def route(body: bytes, remote: str | None, up: list[bool]) -> int:
    """Convenience wrapper: decode the request body and choose a replica."""
    return pick(*affinity_key(body, remote), up)