#!/usr/bin/env python3
"""Fronset cookbook 5: your own keys, how comparison runs, and how to switch it off.

While a task of your own is still being measured, a call to it can run more than
one model: your call is answered, the other models in the task's group answer
the same input, and a judge compares the answers. On your own keys every one of
those model calls runs on your provider account. This cookbook makes two runs of
the same 100 inputs on two new tasks, one with comparison on and one with it
switched off, and prints which model calls ran behind them.

    pip install httpx
    export FRONSET_MANAGEMENT_KEY=llmb_m_...    # the vault and the tasks
    export FRONSET_INFERENCE_KEY=llmb_i_...     # the calls
    export PROVIDER_SLUG=openai                 # optional: vault this provider...
    export PROVIDER_API_KEY=sk-...              # ...with this key (never printed)
    export FRONSET_PORTAL_TOKEN=llmbp_...       # optional: three route details
    python cb5_your_own_keys.py

Later (judging that runs on the batch lane can finish hours after the calls):

    python cb5_your_own_keys.py --recheck record.json   # needs FRONSET_PORTAL_TOKEN

Keys are read from the environment only. Nothing here prints a key, and the
provider key goes to Fronset's vault and nowhere else.

The inputs are 100 synthetic support tickets built below: made-up products and
problems, no real people. Use public or synthetic data on the free plan: see the
plan note on the cookbook page.
"""
import argparse
import datetime as dt
import json
import os
import re
import socket
import sys
import time

import httpx

API = os.environ.get("FRONSET_API_BASE", "https://api.fronset.ai")
PORTAL = os.environ.get("FRONSET_PORTAL_BASE", "https://fronset.ai")
MG = f"{API}/api/inference/v1/management"

# The task: sort a support ticket into one of five categories.
CATEGORIES = ["billing", "delivery", "defect", "account", "other"]
TASK_TEMPLATES = [{"key": "default",
                   "system": "You sort customer support tickets. Answer with the one "
                             "category that fits best: " + ", ".join(CATEGORIES) + ".",
                   "user": "Ticket: {ticket}",
                   "variables": [{"name": "ticket", "required": True,
                                  "description": "The ticket text."}]}]
TASK_OUTPUT_SCHEMA = {"type": "object",
                      "properties": {"category": {"type": "string", "enum": CATEGORIES}},
                      "required": ["category"], "additionalProperties": False}

# Two cheap models from one provider, in two different model families. A
# comparison needs at least two families in the group; one family never compares.
TASK_MODELS = ["gpt-5.4-nano", "gpt-6-luna"]

# 100 synthetic tickets: 10 made-up products x 10 problems.
PRODUCTS = ["Lumen desk lamp", "Brook kettle", "Tarn backpack", "Fennel blender",
            "Corrie headphones", "Holt standing desk", "Marl rain jacket",
            "Pike coffee grinder", "Rill water bottle", "Scree hiking boots"]
PROBLEMS = ["I was charged twice for my {p}.",
            "My {p} was meant to arrive on Monday and it still hasn't come.",
            "The {p} stopped working after three days.",
            "I can't log in to see the order for my {p}.",
            "The parcel with my {p} was left at the wrong address.",
            "The refund for my {p} has not reached my card.",
            "The {p} arrived with a cracked part.",
            "How do I change the email on the account I used to buy the {p}?",
            "Do you sell a spare part for the {p}?",
            "The invoice for my {p} shows the wrong VAT number."]
TICKETS = [q.format(p=p) for q in PROBLEMS for p in PRODUCTS]

CALL_TIMEOUT_S = 900       # a compared call runs several models and a judge


# While a compared call runs, no bytes cross the connection, and a home router,
# NAT or VPN can drop a connection that sits idle for a few minutes. TCP
# keepalives every 30 seconds keep the line in use until the answer comes back.
KEEPALIVE = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
for _name, _value in (("TCP_KEEPIDLE", 30), ("TCP_KEEPALIVE", 30),
                      ("TCP_KEEPINTVL", 15), ("TCP_KEEPCNT", 4)):
    if hasattr(socket, _name):              # Linux has KEEPIDLE, macOS KEEPALIVE
        KEEPALIVE.append((socket.IPPROTO_TCP, getattr(socket, _name), _value))
CALLS = httpx.Client(transport=httpx.HTTPTransport(socket_options=KEEPALIVE),
                     timeout=CALL_TIMEOUT_S)

# The written transcript shows "<account>." for your account number; the
# terminal shows the real slug (see --show-account).
MASKS = []


class Transcript:
    """What one run printed: the page shows `lines`, never the request ids."""

    def __init__(self):
        self.lines, self.record = [], {}

    def say(self, line=""):
        print(line, flush=True)
        for pattern, shown in MASKS:
            line = pattern.sub(shown, line)
        self.lines.append(line)


def mgmt():
    return {"Authorization": f"Bearer {os.environ['FRONSET_MANAGEMENT_KEY']}"}


def inference():
    return {"Authorization": f"Bearer {os.environ['FRONSET_INFERENCE_KEY']}"}


def body_of(r):
    try:
        b = r.json()
    except ValueError:
        return {}
    return b if isinstance(b, dict) else {}


def refused(what, r):
    """Stop with the API's own reason. Never a key."""
    b = body_of(r)
    sys.exit(f"{what}: HTTP {r.status_code}, {b.get('code') or b.get('error_code')}: "
             f"{b.get('error') or b.get('error_message') or b.get('detail')}")


def vault(t):
    """Step 1. Vault one provider key, then read the vault back: metadata only."""
    provider = os.environ.get("PROVIDER_SLUG")
    if os.environ.get("PROVIDER_API_KEY"):
        if not provider:
            sys.exit("set PROVIDER_SLUG to the provider whose key PROVIDER_API_KEY is")
        r = httpx.put(f"{API}/api/inference/v1/vault/providers/{provider}", headers=mgmt(),
                      json={"api_key": os.environ["PROVIDER_API_KEY"], "config": {}}, timeout=60)
        if r.status_code != 200:
            refused("vault refused", r)
        t.say(f"vaulted: {provider} (verified={r.json()['verified']})")
    r = httpx.get(f"{API}/api/inference/v1/vault/providers", headers=mgmt(), timeout=30)
    r.raise_for_status()
    t.say("the vault, as the API returns it (the key itself is never returned):")
    for c in sorted(r.json(), key=lambda c: c["provider_slug"]):
        t.say(f"  {c['provider_slug']:<12} active={c['is_active']}  "
              f"key_hint=…{c['key_hint']}  verified={c['verified']}")
    active = {c["provider_slug"] for c in r.json() if c["is_active"]}
    if not active:
        sys.exit("vault a provider key first (PROVIDER_SLUG + PROVIDER_API_KEY)")
    return active


def ensure_task(t, name, mask):
    """Step 2. Reuse the task if it exists and is active, else create it."""
    r = httpx.get(f"{MG}/account", headers=mgmt(), timeout=30)
    if r.status_code != 200:
        refused("account read refused", r)
    prefix = r.json()["task_prefix"]
    if mask and prefix and not MASKS:
        account = re.escape(prefix.rstrip("."))
        MASKS.append((re.compile(rf"(?<![\w.$]){account}\.(?=[a-z])"), "<account>."))
    r = httpx.get(f"{MG}/capabilities/{prefix}{name}", headers=mgmt(), timeout=30)
    if r.status_code == 200:
        task = r.json()
        if not task["is_active"]:
            sys.exit(f"task {name} exists but is paused; pass a new --tag")
        status = (task.get("maturity") or {}).get("status")
        t.say(f"task: {task['slug']} (existing, {status})")
        return task["slug"]
    if r.status_code != 404:
        refused("task lookup refused", r)
    r = httpx.post(f"{MG}/capabilities/", headers=mgmt(), timeout=60, json={
        "slug": name, "display_name": "Cookbook 5: sort a support ticket",
        "description": "Sort a support ticket into one of five categories.",
        "pricing_mode": "byok", "templates": TASK_TEMPLATES,
        "output_format": "json_schema", "output_schema": TASK_OUTPUT_SCHEMA,
        "model_group": {"name": f"{name[:41]}-models", "models": TASK_MODELS}})
    if r.status_code != 201:
        refused("task create refused", r)
    task = r.json()
    ex = task.get("exploration") or {}
    t.say(f"task: {task['slug']} (created now)")
    t.say(f"  models: {', '.join(sorted(c['model_slug'] for c in task['cells'] if c['is_active']))}")
    fans = ex.get("fans_out") or {}
    t.say(f"  status {ex.get('status')}, compares: native {fans.get('native')}, "
          f"/v1 {fans.get('v1')}; judging: {ex.get('judging')}")
    if ex.get("judging") == "no_judge_keys":
        sys.exit("your vaulted keys can run no judge model, so this task will never "
                 "compare: vault a key for a judge provider (see the page)")
    return task["slug"]


def invoke(slug, workflow_id, ticket, comparison):
    """One call on the native endpoint. `prompt_vars` fills the task's template.
    With comparison off, the call says so in `selection`."""
    body = {"capability_slug": slug, "workflow_id": workflow_id,
            "prompt_vars": {"ticket": ticket}}
    if not comparison:
        body["selection"] = {"fan_out_strategy": "never"}
    return CALLS.post(f"{API}/api/inference/v1/invocations/", headers=inference(), json=body)


def mark_of(c, active):
    """What one model call in a route was for."""
    if c.get("status") == "in_flight":
        return "still running"
    return ("the answer" if c.get("delivered") else
            "an alternative" if c.get("fanout_role") == "shadow" else
            "a repair" if c.get("purpose") == "llm_repair" else
            "on Fronset's key" if c.get("provider_slug") not in active else
            "the judge" if c.get("status") == "success" else c.get("status"))


def get_route(request_id):
    """One route detail, waiting out the portal's rate limit. None if unreadable."""
    for _ in range(5):
        try:
            r = httpx.get(f"{PORTAL}/portal/api/v1/routes/{request_id}",
                          headers={"Authorization": f"Bearer {os.environ['FRONSET_PORTAL_TOKEN']}"},
                          timeout=60)
        except httpx.HTTPError:
            return None
        if r.status_code == 429:
            time.sleep(float(r.headers.get("Retry-After") or 1))
            continue
        if r.status_code == 401:
            sys.exit("the portal token has expired: mint a new one and run --recheck again")
        return r.json() if r.status_code == 200 else None
    return None


def recheck(t, record, active):
    """Some judge calls run on the batch lane and finish later. Read every call's
    route again and count what is behind each run now."""
    for name, label in (("cb5_on", "on"), ("cb5_off", "off")):
        ids = (record.get(name) or {}).get("request_ids") or []
        if not ids:
            continue
        tally, unread = {}, 0
        for request_id in ids:
            route = get_route(request_id)
            if route is None:
                unread += 1
                continue
            for c in route["attempts_detail"]:
                tally[mark_of(c, active)] = tally.get(mark_of(c, active), 0) + 1
            time.sleep(0.25)
        t.say(f"comparison {label}: {len(ids) - unread} of {len(ids)} calls read")
        for mark, calls in sorted(tally.items()):
            t.say(f"  {mark:<17} {calls:>4} model calls")


def route_detail(t, request_id, active):
    """Every model call behind one call (portal token)."""
    # Optional: a route detail that cannot be read never stops the run.
    try:
        r = httpx.get(f"{PORTAL}/portal/api/v1/routes/{request_id}",
                      headers={"Authorization": f"Bearer {os.environ['FRONSET_PORTAL_TOKEN']}"},
                      timeout=60)
    except httpx.HTTPError as e:
        t.say(f"  route detail: skipped ({type(e).__name__})")
        return
    if r.status_code == 401:
        t.say("  route detail: skipped (the portal token has expired)")
        return
    if r.status_code != 200:
        t.say(f"  route detail: skipped (HTTP {r.status_code})")
        return
    route = r.json()
    t.say(f"  route detail: {route['attempts']} model calls, "
          f"${route['route_cost_usd'] or 0:.8f} at list")
    for c in route["attempts_detail"]:
        # A call on a provider you have not vaulted ran on Fronset's key: for
        # example adapting the task's prompt to a model. Of the rest, the one
        # call that is neither the answer, an alternative nor a repair is the judge.
        mark = mark_of(c, active)
        t.say(f"    {c.get('model_slug') or '-':<32} {c.get('provider_slug') or '-':<10} "
              f"${c.get('list_cost_usd') or 0:.8f}  {mark}")


def run(t, slug, tag, comparison, n, details, active):
    """Step 3 or 4: `n` calls, and which model answered each."""
    label = "on" if comparison else "off"
    workflow_id = f"cookbook-5-{tag}-{label}"
    compared = answered = 0
    models = {}
    t.say(f"run: comparison {label}, {n} calls, started "
          f"{dt.datetime.now(dt.timezone.utc):%Y-%m-%d %H:%M} UTC")
    for i, ticket in enumerate(TICKETS[:n]):
        try:
            r = invoke(slug, workflow_id, ticket, comparison)
        except httpx.TransportError as e:
            # No response came back. The call may still have run on the server
            # (and on your key), so check the dashboard before running it again.
            t.record.setdefault("refusals", []).append({"transport": type(e).__name__})
            t.say(f"call {i + 1} failed: no response ({type(e).__name__}); the call may "
                  "still have run, so check the dashboard before running again")
            return False
        b = body_of(r)
        if r.status_code != 200 or b.get("status") != "success":
            t.record.setdefault("refusals", []).append(
                {k: b.get(k) for k in ("error_category", "error_code", "error_message",
                                       "request_id", "call_id")} | {"http": r.status_code})
            t.say(f"call {i + 1} refused: HTTP {r.status_code}, {b.get('error_code')} "
                  f"({b.get('error_category')}): {b.get('error_message')}")
            t.say("stopped: see the error above")
            return False
        answered += 1
        models[b["model_id"]] = models.get(b["model_id"], 0) + 1
        if b.get("panel_version") is not None:
            compared += 1
        t.record.setdefault("request_ids", []).append(b.get("request_id"))
        if details and len(t.record.setdefault("inspected", [])) < 3 and (
                b.get("panel_version") is not None or not comparison):
            t.record["inspected"].append(b.get("request_id"))
            t.say(f"call {i + 1}: answered by {b['model_id']} "
                  f"(${b['cost_usd']:.8f}), category {b['parsed'].get('category')}")
            route_detail(t, b.get("request_id"), active)
    t.say(f"answered: {answered}; by {', '.join(f'{m} {c}' for m, c in sorted(models.items()))}")
    t.say(f"calls that ran a comparison: {compared}")
    return True


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--tag", default=f"{dt.date.today():%Y%m%d}",
                   help="names the two tasks cookbook-5-<tag>-on / -off")
    p.add_argument("--calls", type=int, default=100, help="calls per run (at most 100)")
    p.add_argument("--show-account", action="store_true",
                   help="keep your account number in the written transcript")
    p.add_argument("--page-dir", help="write each run's transcript here as JSON")
    p.add_argument("--record-json", help="write request ids here")
    p.add_argument("--recheck", metavar="RECORD_JSON",
                   help="later: read every call's route again from a record written by a "
                        "full run (needs FRONSET_PORTAL_TOKEN), and make no calls")
    a = p.parse_args()
    if a.recheck:
        if not os.environ.get("FRONSET_PORTAL_TOKEN"):
            sys.exit("--recheck reads route details: export FRONSET_PORTAL_TOKEN first")
        with open(a.recheck) as f:
            record = json.load(f)
        t = Transcript()
        t.say(f"recheck: {dt.datetime.now(dt.timezone.utc):%Y-%m-%d %H:%M} UTC")
        os.environ.pop("PROVIDER_API_KEY", None)  # a recheck never writes the vault
        recheck(t, record, vault(t))
        if a.page_dir:
            with open(os.path.join(a.page_dir, "cb5_recheck.json"), "w") as f:
                json.dump({"run_date": f"{dt.date.today():%Y-%m-%d}", "transcript": t.lines},
                          f, indent=2)
        return
    n = min(a.calls, len(TICKETS))
    start = dt.datetime.now(dt.timezone.utc).date()
    details = bool(os.environ.get("FRONSET_PORTAL_TOKEN"))

    runs = {}
    try:
        setup = runs["cb5_setup"] = Transcript()
        setup.say(f"run: {dt.datetime.now(dt.timezone.utc):%Y-%m-%d %H:%M} UTC against {API}")
        active = vault(setup)
        slug_on = ensure_task(setup, f"cookbook-5-{a.tag}-on", not a.show_account)
        slug_off = ensure_task(setup, f"cookbook-5-{a.tag}-off", not a.show_account)
        if run(runs.setdefault("cb5_on", Transcript()), slug_on, a.tag, True, n, details, active):
            run(runs.setdefault("cb5_off", Transcript()), slug_off, a.tag, False, n, details,
                active)
    finally:
        # Written even if the run stops part-way, so what ran is not lost.
        if a.page_dir:
            for name, t in runs.items():
                with open(os.path.join(a.page_dir, f"{name}.json"), "w") as f:
                    json.dump({"run_date": f"{start:%Y-%m-%d}", "transcript": t.lines}, f, indent=2)
        if a.record_json:
            with open(a.record_json, "w") as f:
                json.dump({n_: t.record for n_, t in runs.items()}, f, indent=2)


if __name__ == "__main__":
    main()