#!/usr/bin/env python3
"""Fronset cookbook 1: which model answered, and why.

One call on a catalogue task (managed: we pay the provider, you pay us), and
one call on a task of your own (own keys: your provider bills you). The own-key
half creates its task first, through the management API, and calls the slug the
create returns: /v1 only calls tasks that exist. Both calls are in template mode:
the task's template builds the prompt from the variables you send, so `messages`
is empty. For each call, every value that says which model answered and what that
cost.

    pip install openai          # httpx comes with it
    export FRONSET_INFERENCE_KEY=llmb_i_...     # the /v1 key
    export FRONSET_MANAGEMENT_KEY=llmb_m_...    # the vault, task and usage key
    export FRONSET_PORTAL_TOKEN=llmbp_...       # optional: route detail
    export PROVIDER_SLUG=openai                 # optional: vault this provider...
    export PROVIDER_API_KEY=sk-...              # ...with this key (never printed)
    python cb1_which_model_answered.py --path both

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 text sent is the opening of the Gettysburg Address (1863, public domain).
Use public 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 httpx
import openai
from openai import OpenAI

API = os.environ.get("FRONSET_API_BASE", "https://api.fronset.ai")
PORTAL = os.environ.get("FRONSET_PORTAL_BASE", "https://fronset.ai")
CATALOGUE_TASK = "structured-content-summarization"

TEXT = (
    "Four score and seven years ago our fathers brought forth on this continent, "
    "a new nation, conceived in Liberty, and dedicated to the proposition that all "
    "men are created equal. Now we are engaged in a great civil war, testing whether "
    "that nation, or any nation so conceived and so dedicated, can long endure. We are "
    "met on a great battle-field of that war. We have come to dedicate a portion of "
    "that field, as a final resting place for those who here gave their lives that "
    "that nation might live. It is altogether fitting and proper that we should do this."
)

# The per-call fee card (Pricing Terms): the first 10,000 calls a month carry no
# fee, then $2.00 / $1.50 / $1.00 per 1,000, graduated.
FEE_BANDS = [(10_000, 0.0), (50_000, 2.00), (250_000, 1.50), (None, 1.00)]

# The x-llmbench-* values the body mirrors under "llmbench" (prefix stripped).
MIRRORED = ["capability", "model", "provider", "service-tier", "cost-usd",
            "band", "phase"]

# The own-key task's output: one field.
TASK_OUTPUT_SCHEMA = {"type": "object", "properties": {"summary": {"type": "string"}},
                      "required": ["summary"], "additionalProperties": False}

# The own-key task's template: its two prompt texts and the one variable they
# read. A call to the task sends the variable's value; the text stays here.
TASK_TEMPLATES = [{"key": "default",
                   "system": "Summarise the text in one sentence.",
                   "user": "{text}",
                   "variables": [{"name": "text", "required": True,
                                  "description": "The text to summarise."}]}]

# The models the own-key task may run: two, from the default provider above,
# so the pool is the one you chose (the create response tells you to narrow it).
TASK_MODELS = ["gpt-6-sol", "gpt-6-luna"]

# The catalogue task's template variables (GET /v1/models/{task} lists them).
CATALOGUE_VARS = {
    "content_text": TEXT,
    "content_title": "The Gettysburg Address (opening)",
    "content_type": "speech",
    "content_url": "",
    "summary_structure": "## Summary\nOne sentence on what the speech says.",
}

# A call to a task that is still exploring can run several models and a judge, so
# it can take much longer than a plain call. A retried call runs again from the
# start and its model usage is billed again, so the client never retries.
CALL_TIMEOUT_S = 900

# While the 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))


def keepalive_http_client():
    """The OpenAI SDK's own HTTP client, with keepalives on its sockets."""
    http = sys.modules[openai.DefaultHttpxClient.__mro__[1].__module__.split(".")[0]]
    return openai.DefaultHttpxClient(transport=http.HTTPTransport(socket_options=KEEPALIVE))

# The written transcript hides your account number (see --show-account).
MASK_ACCOUNT = True


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

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

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


def per_call_fee(calls_before, calls):
    """Fee for `calls` more calls this month, after `calls_before` already made."""
    fee, lower = 0.0, 0
    for upper, per_1k in FEE_BANDS:
        top = upper if upper is not None else calls_before + calls
        in_band = max(0, min(calls_before + calls, top) - max(calls_before, lower))
        fee += in_band * per_1k / 1000
        lower = top
    return fee


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


def refused(what, r):
    """Stop with the API's own reason, its `code` and `error`. Never a key."""
    try:
        body = r.json()
    except ValueError:
        body = {}
    if not isinstance(body, dict):
        body = {}
    sys.exit(f"{what}: HTTP {r.status_code}, code {body.get('code')}: {body.get('error')}")


def ensure_task(name, schema=True):
    """Step 3. Your own task: reuse it if it exists and is active, else create it.

    Returns (slug, created). The slug is the one the API returns,
    `<account id>.<name>`, and it is what /v1 calls. A paused task is never
    reused: a call to it would 404."""
    mg = f"{API}/api/inference/v1/management"
    r = httpx.get(f"{mg}/account", headers=mgmt_headers(), timeout=30)
    if r.status_code != 200:
        refused("account read refused", r)
    prefix = r.json()["task_prefix"]            # your account number and a dot, e.g. "3."
    r = httpx.get(f"{mg}/capabilities/{prefix}{name}", headers=mgmt_headers(), timeout=30)
    if r.status_code == 200:
        task = r.json()
        if not task["is_active"]:
            sys.exit(f"task {task['slug']} exists but is paused; pass a new --task name")
        return task["slug"], False
    if r.status_code != 404:
        refused("task lookup refused", r)
    body = {"slug": name, "display_name": "Cookbook 1: one-sentence summary",
            "description": "Summarise a text in one sentence.", "pricing_mode": "byok",
            "templates": TASK_TEMPLATES,
            # A new group, named after the task: a second task gets its own
            # (a group name is taken once). Group names are at most 48 characters.
            "model_group": {"name": f"{name[:41]}-models", "models": TASK_MODELS}}
    if schema:
        body["output_format"] = "json_schema"
        body["output_schema"] = TASK_OUTPUT_SCHEMA
    else:
        body["output_format"] = "text"
    r = httpx.post(f"{mg}/capabilities/", headers=mgmt_headers(), json=body, timeout=60)
    if r.status_code != 201:
        refused("task create refused", r)
    return r.json()["slug"], True


def calls_this_month(today):
    """Delivered calls so far this month, for the per-call fee arithmetic."""
    r = httpx.get(f"{API}/api/inference/v1/usage", headers=mgmt_headers(), timeout=30,
                  params={"from": today.replace(day=1).isoformat(),
                          "to": today.isoformat(), "group_by": "role"})
    r.raise_for_status()
    return sum(row["calls"] for row in r.json()["rows"] if row["key"] == "delivered")


def call_and_read(client, t, task, prompt_vars):
    """Make one call in template mode and print which model answered, from headers
    and body. The task's template renders `prompt_vars`; `messages` stays empty."""
    try:
        raw = client.chat.completions.with_raw_response.create(
            model=task,
            messages=[],
            extra_body={"llmbench": {"prompt_vars": prompt_vars}},
        )
    except openai.APIStatusError as e:
        h = e.response.headers
        t.record["diagnostic"] = h.get("x-llmbench-diagnostic")
        t.say(f"call failed: HTTP {e.status_code}, {h.get('x-llmbench-error-code')} "
              f"(whose problem: {h.get('x-llmbench-error-class')})")
        return None
    except openai.APIConnectionError as e:
        # No response came back. The call may still have run (and been billed)
        # on the server, so check usage before running it again.
        t.say(f"call failed: no response ({type(e).__name__}); the call may still have run, "
              "so check GET /api/inference/v1/usage before calling again")
        return None
    h = raw.headers
    completion = raw.parse()
    t.record["request_id"] = h.get("x-llmbench-request-id")
    t.record["call_id"] = h.get("x-llmbench-call-id")

    t.say(f"task: {h.get('x-llmbench-capability')}")
    t.say(f"answer: {completion.choices[0].message.content.strip()}")
    t.say()
    t.say("from the headers")
    t.say(f"  model     {h.get('x-llmbench-model')}")
    t.say(f"  provider  {h.get('x-llmbench-provider')}")
    t.say(f"  cost-usd  {h.get('x-llmbench-cost-usd')}   (the answering model's list cost)")
    t.say(f"  band      {h.get('x-llmbench-band', '(absent)')}")
    t.say(f"  phase     {h.get('x-llmbench-phase', '(absent)')}")

    # Step 4. The same values, mirrored into the body under "llmbench".
    body = (completion.model_extra or {}).get("llmbench") or {}
    same = all(body.get(k) == h.get(f"x-llmbench-{k}") for k in MIRRORED)
    t.say(f"body 'llmbench' object matches the headers: {'yes' if same else 'NO'}")
    return h, completion.usage


def route_detail(t, request_id):
    """Step 5. Every model call behind your one call (portal token)."""
    portal = os.environ.get("FRONSET_PORTAL_TOKEN")
    if not (portal and request_id):
        t.say("route detail: skipped (set FRONSET_PORTAL_TOKEN)")
        return None
    r = httpx.get(f"{PORTAL}/portal/api/v1/routes/{request_id}",
                  headers={"Authorization": f"Bearer {portal}"}, timeout=30)
    if r.status_code == 401:
        t.say("route detail: skipped (the portal token has expired; mint a new one)")
        return None
    r.raise_for_status()
    route = r.json()
    t.say()
    t.say(f"route detail ({dt.datetime.now(dt.timezone.utc):%H:%M:%S} UTC)")
    t.say(f"  attempts        {route['attempts']}")
    t.say(f"  route_cost_usd  {route['route_cost_usd']:.8f}")
    for c in route["attempts_detail"]:
        mark = "delivered" if c["delivered"] else (
            c.get("purpose") or c.get("fanout_role") or c["status"])
        t.say(f"    {c['model_slug'] or '-':<32} {c['provider_slug'] or '-':<12} "
              f"{c['list_cost_usd'] or 0:.8f}  {mark}")
    return route


def managed(client, today):
    """Step 2. A catalogue task: the schema comes with the task, so no response_format."""
    t = Transcript()
    t.say(f"run: {dt.datetime.now(dt.timezone.utc):%Y-%m-%d %H:%M} UTC against {API}/v1")
    before = calls_this_month(today)
    got = call_and_read(client, t, CATALOGUE_TASK, CATALOGUE_VARS)
    if got is None:
        return t
    h, usage = got
    route_detail(t, t.record.get("request_id"))

    # The rate this call is priced at. The task names its selection policy; the
    # policy's quality bar and evidence floor are the shape the rate is keyed on.
    mg = f"{API}/api/inference/v1/management"
    r = httpx.get(f"{mg}/capabilities/{CATALOGUE_TASK}", headers=mgmt_headers(), timeout=30)
    if r.status_code != 200:
        refused("task read refused", r)
    task = r.json()
    params = {"lane": "sync", "eligibility_floor": "MEDIUM"}
    if task.get("selection_policy_slug"):
        r = httpx.get(f"{mg}/selection-policies/{task['selection_policy_slug']}",
                      headers=mgmt_headers(), timeout=30)
        if r.status_code != 200:
            refused("selection policy read refused", r)
        policy = r.json()
        params["quality_step"] = f"{policy['quality_threshold']:.2f}"
        params["eligibility_floor"] = policy["min_confidence"] or "MEDIUM"
    if task.get("model_group"):
        params["model_group"] = task["model_group"]
    r = httpx.get(f"{API}/api/billing/v1/rates/{CATALOGUE_TASK}/quote", headers=mgmt_headers(),
                  timeout=30, params=params)
    # 404: the task has no published rate at all. That is an answer, not an error.
    if r.status_code == 404:
        q = {"priced": False}
    elif r.status_code != 200:
        refused("rate quote refused", r)
    else:
        q = r.json()
    res = q.get("resolved") or {}
    t.say()
    t.say("what this run cost us")
    if q["priced"] and res.get("input") and res.get("output"):
        rate_in = float(res["input"]["rate"])
        rate_out = float(res["output"]["rate"])
        charge = (usage.prompt_tokens * rate_in + usage.completion_tokens * rate_out) / 1e6
        t.say(f"  published rate       ${rate_in:.3f} in / ${rate_out:.3f} out per 1M tokens "
              f"(quality {res['quality_step']}, floor {res['eligibility_floor']}, since "
              f"{res['output']['effective_from'][:10]})")
        t.say(f"  model usage          ${charge:.8f}  ({usage.prompt_tokens} in, "
              f"{usage.completion_tokens} out, at the published rate)")
    elif q["priced"]:
        t.say("  published rate       only one direction is published for this call's shape; "
              "model usage is not worked out here")
    else:
        t.say("  published rate       NONE in force for this call's shape: this call is "
              "billed at cost (list price, no markup)")
    t.say(f"  per-call fee         ${per_call_fee(before, 1):.4f}  (call {before + 1:,} "
          f"this month; the first 10,000 carry no fee)")
    t.say("  credit fee           5.5% of credit bought (minimum $0.80), paid when you buy; none bought for this run")
    return t


def own_key(client, today, task, schema=True):
    """Step 3. Your own task, on your own provider key: vault, create, call."""
    t = Transcript()
    t.say(f"run: {dt.datetime.now(dt.timezone.utc):%Y-%m-%d %H:%M} UTC against {API}/v1")

    # Vault one provider key with the management key.
    provider = os.environ.get("PROVIDER_SLUG")
    if os.environ.get("PROVIDER_API_KEY"):
        if not provider:
            t.say("stopped: set PROVIDER_SLUG to the provider whose key PROVIDER_API_KEY is")
            return t
        r = httpx.put(f"{API}/api/inference/v1/vault/providers/{provider}", headers=mgmt_headers(),
                      json={"api_key": os.environ["PROVIDER_API_KEY"], "config": {}},
                      timeout=60)
        r.raise_for_status()
        t.say(f"vaulted: {provider} (verified={r.json()['verified']})")
    r = httpx.get(f"{API}/api/inference/v1/vault/providers", headers=mgmt_headers(), timeout=30)
    r.raise_for_status()
    active = sorted(c["provider_slug"] for c in r.json() if c["is_active"])
    t.say(f"providers in the vault: {', '.join(active) or 'none'}")
    if not active:
        t.say("stopped: vault a provider key first (PROVIDER_SLUG + PROVIDER_API_KEY)")
        return t

    # Create the task (or reuse it), with the management key. /v1 calls it by
    # the slug the API returns, never by the bare name.
    slug, created = ensure_task(task, schema)
    if MASK_ACCOUNT and "." in slug:
        # The page transcript shows "<account>." for your account number; the
        # terminal shows the real slug.
        account = re.escape(slug.split(".", 1)[0])
        t.masks.append((re.compile(rf"(?<![\w.$]){account}\.(?=[a-z])"), "<account>."))
    t.say(f"task: {slug} ({'created now' if created else 'existing'})")

    before = calls_this_month(today)
    t.say()
    got = call_and_read(client, t, slug, {"text": TEXT})
    if got is None:
        return t
    h, _ = got
    route = route_detail(t, t.record.get("request_id"))

    t.say()
    t.say("what this run cost us")
    if route:
        # Calls on a provider in your vault ran on your key; any other call in the
        # route (for example adapting the prompt to a new model) ran on ours.
        yours = sum(c["list_cost_usd"] or 0 for c in route["attempts_detail"]
                    if c["provider_slug"] in active)
        ours = (route["route_cost_usd"] or 0) - yours
        t.say(f"  model usage at list  ${yours:.8f}  (billed by the provider to the vaulted key)")
        if ours > 0:
            t.say(f"  on Fronset's keys    ${ours:.8f}  (not billed to you)")
    else:
        t.say(f"  model usage at list  ${float(h.get('x-llmbench-cost-usd') or 0):.8f}  "
              "(the answering model only; the route detail has every call)")
    t.say(f"  per-call fee         ${per_call_fee(before, 1):.4f}  (call {before + 1:,} "
          f"this month; the first 10,000 carry no fee)")
    t.say("  credit fee           $0.00  (no credit bought on the own-key path)")
    return t


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--path", choices=["managed", "own-key", "both"], default="both")
    p.add_argument("--task", default=f"cookbook-1-summary-{dt.date.today():%Y%m%d}",
                   help="the name of the own-key task to create, or reuse if it exists and "
                        "is active; the call uses the slug the create returned")
    p.add_argument("--no-schema", action="store_true",
                   help="create the own-key task without an output schema (its answer is "
                        "then plain text)")
    p.add_argument("--page-dir", help="write each path's transcript here as JSON")
    p.add_argument("--record-json", help="write request ids here")
    p.add_argument("--show-account", action="store_true",
                   help="keep your account number in the written transcript")
    a = p.parse_args()
    global MASK_ACCOUNT
    MASK_ACCOUNT = not a.show_account
    today = dt.datetime.now(dt.timezone.utc).date()

    # Step 1. Point the OpenAI SDK at Fronset, with the inference key.
    client = OpenAI(base_url=f"{API}/v1", api_key=os.environ["FRONSET_INFERENCE_KEY"],
                    timeout=CALL_TIMEOUT_S, max_retries=0,
                    http_client=keepalive_http_client())

    runs = {}
    if a.path in ("managed", "both"):
        print("== managed: a catalogue task ==")
        runs["cb1_managed"] = managed(client, today)
    if a.path in ("own-key", "both"):
        print("\n== own keys: a task you create ==")
        runs["cb1_own_key"] = own_key(client, today, a.task, schema=not a.no_schema)

    for name, t in runs.items():
        if a.page_dir:
            with open(os.path.join(a.page_dir, f"{name}.json"), "w") as f:
                json.dump({"run_date": f"{today:%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()