"""
fable5-delegate: escalate a hard task to Claude Fable 5, fall back to Opus 4.8.

Fable 5 is Anthropic's most capable widely-released model (GA 2026-06-09): 1M
context, priced above the Opus tier, and slower. You don't want it for everything,
only for the hardest ~10% of a job, with a safe, automatic path back to
Claude Opus 4.8 for the rest.

This module is that escalation boundary. Give it a well-specified hard task; it
calls `claude-fable-5`, falls back to `claude-opus-4-8` when Fable's safety
classifiers decline, and hands back the answer plus an honest readout of which
model actually served it and what it cost.

Two fallback mechanisms, both handled here:
  1. SERVER-SIDE (Anthropic's recommended default): pass the
     `server-side-fallback-2026-06-01` beta and a `fallbacks` list; on a policy
     decline the API re-serves the SAME request on the fallback model inside one
     call, with credit-style repricing.
  2. CLIENT-SIDE (the resilient floor): call Fable plainly, and if it returns a
     refusal, re-send the identical request to Opus 4.8 yourself. Uses only
     stable, GA API surface, so it works even if the server-side beta is
     unavailable in your account or its wire contract changes.
`delegate()` tries the server-side path first and DEGRADES to the client-side
path automatically if the server rejects the fallback beta. `fallback_mode` on
the result tells you which mechanism was in effect.

Why the other specifics (all verified against Anthropic's model docs):
  - Thinking is ALWAYS ON for Fable 5, so you omit the `thinking` parameter
    entirely. An explicit `{"type": "disabled"}` returns a 400.
  - `temperature` / `top_p` / `top_k` are removed on Fable 5. Steer with the
    prompt and the `effort` control instead. Sending them returns a 400.
  - Depth is controlled by `output_config.effort` (low | medium | high | xhigh | max).
  - Refusals arrive as HTTP 200 with `stop_reason == "refusal"`, so you must check
    stop_reason BEFORE reading content, or you'll index into an empty list.
  - Fable 5 requires 30-day data retention; zero-data-retention orgs get a 400.

Requires: anthropic>=0.116, ANTHROPIC_API_KEY in the environment.
Usage:     python fable5_delegate.py "your hard task here"
Self-test: python fable5_delegate.py --self-test    (offline; no API key needed)
Dry run:   python fable5_delegate.py --dry-run "task"  (prints the request, no call)
"""

from __future__ import annotations

import os
import sys
from dataclasses import dataclass, field
from typing import Any, Optional

FABLE = "claude-fable-5"
FALLBACK = "claude-opus-4-8"
FALLBACK_BETA = "server-side-fallback-2026-06-01"

# Published list prices, USD per million tokens (input, output).
PRICES = {
    "claude-fable-5": (10.0, 50.0),
    "claude-opus-4-8": (5.0, 25.0),
}

# Streaming is required above this output size or the request can hit an HTTP
# timeout before the (slower) Fable turn finishes.
STREAM_ABOVE_MAX_TOKENS = 16_000


@dataclass
class DelegateResult:
    text: str
    served_by: str          # the model that actually produced the answer
    requested: str          # the model we asked for (always Fable here)
    fell_back: bool         # did a fallback model serve this turn?
    fallback_mode: str      # "server-side" | "client-side": which mechanism was armed
    refused: bool           # did the WHOLE chain refuse?
    refusal_category: Optional[str]
    stop_reason: Optional[str]
    switch_points: list[tuple[str, str]]  # (from_model, to_model) per fallback hop
    input_tokens: int
    output_tokens: int
    est_cost_usd: float
    raw: Any = field(repr=False, default=None)

    def summary(self) -> str:
        if self.refused:
            return (f"REFUSED by the full chain (category={self.refusal_category}). "
                    f"No answer produced.")
        note = "" if not self.fell_back else f"  [fell back from Fable 5 via {self.fallback_mode}]"
        return (f"served_by={self.served_by}{note}  "
                f"in={self.input_tokens} out={self.output_tokens} "
                f"~${self.est_cost_usd:.4f}")


def _estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    inp, out = PRICES.get(model, PRICES[FALLBACK])
    return (input_tokens / 1_000_000) * inp + (output_tokens / 1_000_000) * out


def build_request(
    prompt: str,
    *,
    model: str = FABLE,
    effort: str = "high",
    max_tokens: int = 16_000,
    system: Optional[str] = None,
    fallback_model: str = FALLBACK,
    server_fallback: bool = True,
) -> dict[str, Any]:
    """Build the exact kwargs for client.beta.messages.create / .stream.

    Pure function, no network. `parse_result` consumes what these produce.
    With server_fallback=False the fallback beta + `fallbacks` body are omitted,
    which is the plain, GA request used by the client-side fallback path.
    """
    if effort not in ("low", "medium", "high", "xhigh", "max"):
        raise ValueError(f"effort must be low|medium|high|xhigh|max, got {effort!r}")

    kwargs: dict[str, Any] = {
        "model": model,
        "max_tokens": max_tokens,
        "messages": [{"role": "user", "content": prompt}],
        # Depth/cost control. NO thinking param (always-on on Fable; disabling 400s).
        # NO temperature/top_p/top_k (removed on Fable; sending them 400s).
        "output_config": {"effort": effort},
    }
    if system:
        kwargs["system"] = system
    if server_fallback:
        # Opt into transparent recovery so a classifier false-positive doesn't
        # just fail the request. extra_body carries the fallbacks list into the
        # request body regardless of whether the installed SDK types it yet.
        kwargs["betas"] = [FALLBACK_BETA]
        kwargs["extra_body"] = {"fallbacks": [{"model": fallback_model}]}
    return kwargs


def parse_result(msg: Any, requested: str = FABLE) -> DelegateResult:
    """Turn a (Beta)Message into a DelegateResult. Robust to missing beta fields."""
    stop_reason = getattr(msg, "stop_reason", None)
    served_by = getattr(msg, "model", requested) or requested

    # A top-level refusal means Fable AND the fallback both declined.
    refused = stop_reason == "refusal"
    refusal_category = None
    if refused:
        details = getattr(msg, "stop_details", None)
        refusal_category = getattr(details, "category", None)

    # Text = concatenation of text blocks. Fallback blocks mark switch points.
    text_parts: list[str] = []
    switch_points: list[tuple[str, str]] = []
    for block in getattr(msg, "content", None) or []:
        btype = getattr(block, "type", None)
        if btype == "text":
            text_parts.append(getattr(block, "text", "") or "")
        elif btype == "fallback":
            frm = getattr(getattr(block, "from_", None), "model", None) or \
                  getattr(getattr(block, "from", None), "model", None)
            to = getattr(getattr(block, "to", None), "model", None)
            if frm and to:
                switch_points.append((frm, to))

    # Served-by signal: a `fallback_message` in usage.iterations covers sticky
    # turns that carry no fallback block. Pair with stop_reason (a fallback model
    # can itself refuse).
    usage = getattr(msg, "usage", None)
    iterations = getattr(usage, "iterations", None) or []
    fallback_ran = any(getattr(it, "type", None) == "fallback_message" for it in iterations)
    fell_back = bool(switch_points) or (fallback_ran and served_by != requested)

    input_tokens = int(getattr(usage, "input_tokens", 0) or 0)
    output_tokens = int(getattr(usage, "output_tokens", 0) or 0)

    return DelegateResult(
        text="".join(text_parts),
        served_by=served_by,
        requested=requested,
        fell_back=fell_back,
        fallback_mode="server-side",
        refused=refused,
        refusal_category=refusal_category,
        stop_reason=stop_reason,
        switch_points=switch_points,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        est_cost_usd=_estimate_cost(served_by, input_tokens, output_tokens),
        raw=msg,
    )


def _call(client: Any, kwargs: dict[str, Any], max_tokens: int) -> Any:
    """Execute one request, streaming when the output budget is large."""
    if max_tokens > STREAM_ABOVE_MAX_TOKENS:
        with client.beta.messages.stream(**kwargs) as stream:
            return stream.get_final_message()
    return client.beta.messages.create(**kwargs)


def _looks_like_fallback_beta_issue(err_text: str) -> bool:
    t = err_text.lower()
    return any(k in t for k in ("fallback", "server-side-fallback", "beta"))


def _delegate_client_side(
    client: Any, prompt: str, *, effort: str, max_tokens: int,
    system: Optional[str], fallback_model: str,
) -> DelegateResult:
    """Resilient path: plain Fable call, then re-send to Opus 4.8 on a refusal.

    Uses only stable, GA API surface (no beta), so it works regardless of the
    server-side fallback's availability.
    """
    fable_kwargs = build_request(
        prompt, model=FABLE, effort=effort, max_tokens=max_tokens,
        system=system, server_fallback=False,
    )
    r = parse_result(_call(client, fable_kwargs, max_tokens))
    r.fallback_mode = "client-side"
    if not r.refused:
        return r  # Fable answered; no fallback needed.

    # Fable declined. Re-send the identical request to the fallback model.
    fb_kwargs = dict(fable_kwargs, model=fallback_model)
    r2 = parse_result(_call(client, fb_kwargs, max_tokens), requested=FABLE)
    r2.fallback_mode = "client-side"
    if not r2.refused and r2.served_by != FABLE:
        r2.fell_back = True
        r2.switch_points = [(FABLE, r2.served_by)]
    return r2


def delegate(
    prompt: str,
    *,
    effort: str = "high",
    max_tokens: int = 16_000,
    system: Optional[str] = None,
    fallback_model: str = FALLBACK,
    use_server_fallback: bool = True,
    client: Any = None,
) -> DelegateResult:
    """Escalate `prompt` to Fable 5 with a transparent Opus-4.8 fallback.

    Tries Anthropic's server-side fallback first (one round trip, credit
    repricing). If the endpoint rejects that beta, degrades automatically to the
    client-side refusal-retry, which uses only stable API. Streams automatically
    when max_tokens is large (Fable turns can run minutes).
    """
    import anthropic  # imported here so --self-test / import stays dependency-light

    if client is None:
        client = anthropic.Anthropic()  # resolves ANTHROPIC_API_KEY or an `ant` profile

    if use_server_fallback:
        try:
            kwargs = build_request(
                prompt, model=FABLE, effort=effort, max_tokens=max_tokens,
                system=system, fallback_model=fallback_model, server_fallback=True,
            )
            r = parse_result(_call(client, kwargs, max_tokens))
            r.fallback_mode = "server-side"
            return r
        except anthropic.BadRequestError as e:
            text = str(e).lower()
            if "retention" in text:
                raise RuntimeError(
                    "Claude Fable 5 requires 30-day data retention; this organization's "
                    "retention setting is too strict. Enable 30-day retention or route "
                    "this task to claude-opus-4-8 directly."
                ) from e
            if not _looks_like_fallback_beta_issue(text):
                raise  # a real request error, not the fallback beta; surface it
            # Server-side fallback unavailable here; fall through to client-side.

    return _delegate_client_side(
        client, prompt, effort=effort, max_tokens=max_tokens,
        system=system, fallback_model=fallback_model,
    )


# --------------------------------------------------------------------------- #
# Offline self-test: exercises build_request + parse_result against synthetic
# response objects for the real scenarios. No network, no API key.
# --------------------------------------------------------------------------- #
def _self_test() -> int:
    from types import SimpleNamespace as NS

    ok = True

    def check(name: str, cond: bool) -> None:
        nonlocal ok
        print(f"  [{'PASS' if cond else 'FAIL'}] {name}")
        ok = ok and cond

    # 1. server-side request shape
    req = build_request("hard task", effort="xhigh", max_tokens=20000, system="be terse")
    check("model is Fable", req["model"] == FABLE)
    check("fallback beta present", FALLBACK_BETA in req["betas"])
    check("fallbacks in extra_body", req["extra_body"]["fallbacks"] == [{"model": FALLBACK}])
    check("effort under output_config", req["output_config"]["effort"] == "xhigh")
    check("NO thinking param", "thinking" not in req)
    check("NO temperature", "temperature" not in req and "top_p" not in req)
    check("system carried through", req.get("system") == "be terse")
    try:
        build_request("x", effort="bogus")
        check("rejects bad effort", False)
    except ValueError:
        check("rejects bad effort", True)

    # 1b. client-side (plain) request shape
    plain = build_request("x", server_fallback=False)
    check("client-side omits betas", "betas" not in plain and "extra_body" not in plain)
    check("client-side keeps effort/no-thinking",
          plain["output_config"]["effort"] == "high" and "thinking" not in plain)

    # 2. Fable served normally
    fable_ok = NS(
        model="claude-fable-5", stop_reason="end_turn", stop_details=None,
        content=[NS(type="text", text="Fable answered.")],
        usage=NS(input_tokens=1000, output_tokens=500, iterations=[]),
    )
    r = parse_result(fable_ok)
    check("fable text parsed", r.text == "Fable answered.")
    check("served_by fable", r.served_by == "claude-fable-5")
    check("did not fall back", r.fell_back is False)
    check("not refused", r.refused is False)
    check("cost = fable rate", abs(r.est_cost_usd - (0.001 * 10 + 0.0005 * 50)) < 1e-9)

    # 3. Fable refused, Opus 4.8 served (server-side: fallback block + iteration)
    fell_back = NS(
        model="claude-opus-4-8", stop_reason="end_turn", stop_details=None,
        content=[
            NS(type="fallback", from_=NS(model="claude-fable-5"), to=NS(model="claude-opus-4-8")),
            NS(type="text", text="Opus rescued it."),
        ],
        usage=NS(input_tokens=1000, output_tokens=400,
                 iterations=[NS(type="fallback_message")]),
    )
    r = parse_result(fell_back)
    check("fallback text parsed", r.text == "Opus rescued it.")
    check("served_by opus", r.served_by == "claude-opus-4-8")
    check("fell_back true", r.fell_back is True)
    check("switch point recorded", r.switch_points == [("claude-fable-5", "claude-opus-4-8")])
    check("cost = opus rate", abs(r.est_cost_usd - (0.001 * 5 + 0.0004 * 25)) < 1e-9)

    # 4. Whole chain refused
    refused = NS(
        model="claude-fable-5", stop_reason="refusal",
        stop_details=NS(category="cyber", explanation="policy"),
        content=[], usage=NS(input_tokens=0, output_tokens=0, iterations=[]),
    )
    r = parse_result(refused)
    check("refused true", r.refused is True)
    check("category surfaced", r.refusal_category == "cyber")
    check("empty text, no crash", r.text == "")

    # 5. client-side degrade path (no network): a fake client that 400s the beta,
    #    serves Fable-refusal on the plain call, then Opus on the retry.
    class _FakeMessages:
        def __init__(self):
            self.calls = []
        def create(self, **kw):
            self.calls.append(kw)
            if "betas" in kw:  # server-side attempt
                import anthropic
                raise anthropic.BadRequestError(
                    message="unknown anthropic-beta: server-side-fallback-2026-06-01",
                    response=NS(status_code=400, headers={}, request=None), body=None,
                )
            if kw["model"] == FABLE:  # plain Fable -> refuse
                return NS(model=FABLE, stop_reason="refusal",
                          stop_details=NS(category="cyber", explanation="x"),
                          content=[], usage=NS(input_tokens=10, output_tokens=0, iterations=[]))
            return NS(model=kw["model"], stop_reason="end_turn", stop_details=None,
                      content=[NS(type="text", text="client-side Opus answer")],
                      usage=NS(input_tokens=10, output_tokens=20, iterations=[]))
    class _FakeBeta:
        def __init__(self): self.messages = _FakeMessages()
    class _FakeClient:
        def __init__(self): self.beta = _FakeBeta()
    fake = _FakeClient()
    r = delegate("adjacent task", client=fake)
    check("degrade -> client-side mode", r.fallback_mode == "client-side")
    check("degrade served by Opus", r.served_by == "claude-opus-4-8")
    check("degrade fell_back true", r.fell_back is True and not r.refused)
    check("degrade text correct", r.text == "client-side Opus answer")

    print("\nself-test:", "ALL PASS" if ok else "FAILURES ABOVE")
    return 0 if ok else 1


def _main(argv: list[str]) -> int:
    args = [a for a in argv if not a.startswith("--")]
    flags = {a for a in argv if a.startswith("--")}

    if "--self-test" in flags:
        return _self_test()

    if not args:
        print(__doc__)
        return 0

    prompt = args[0]

    if "--dry-run" in flags:
        import json
        print(json.dumps(build_request(prompt), indent=2))
        return 0

    if not os.environ.get("ANTHROPIC_API_KEY"):
        print("Set ANTHROPIC_API_KEY to run a live delegation "
              "(or use --dry-run / --self-test).", file=sys.stderr)
        return 2

    result = delegate(prompt)
    print(result.text)
    print("\n---", result.summary(), file=sys.stderr)
    return 1 if result.refused else 0


if __name__ == "__main__":
    raise SystemExit(_main(sys.argv[1:]))
