#!/usr/bin/env python3
"""Fronset cookbook 4: define a task of your own, call it, and see its progress.

Creates one task of your own (or reuses it when it already exists), reads how
it starts, makes two calls to it on /v1/chat/completions, and reads the task's
maturity before and after. It does not run the task until it matures.

    pip install openai          # httpx comes with it
    export FRONSET_MANAGEMENT_KEY=llmb_m_...    # create and read the task
    export FRONSET_INFERENCE_KEY=llmb_i_...     # the two calls
    python cb4_define_your_own_task.py

The task runs on your own provider keys (pricing_mode "byok"), so vault a key
for the providers of its model group first (Cookbook 5, step 1).

Keys are read from the environment only. Nothing here prints a key. The inputs
are two made-up emails, with no real people in them.
"""
import argparse
import datetime as dt
import json
import os
import sys

import httpx
import openai
from openai import OpenAI

API = os.environ.get("FRONSET_API_BASE", "https://api.fronset.ai")
MGMT = f"{API}/api/inference/v1/management"
NAME = "cookbook-4-action-items"

# The task, as it is created. Two variables, a JSON output, two models from two
# model families on one provider.
TASK = {
    "slug": NAME,
    "display_name": "Cookbook 4: the action an email asks for",
    "description": "Read one email and return the action it asks the reader to "
                   "take, and the date it is due by, if the email gives one.",
    "pricing_mode": "byok",
    "templates": [{
        "key": "default",
        "system": "You read business emails. Return the one action the email asks "
                  "the reader to take, in at most twelve words, and the date it is "
                  "due by as YYYY-MM-DD, or null when the email gives no date.",
        "user": "Subject: {subject}\n\n{body}",
        "variables": [
            {"name": "subject", "required": True, "description": "The email's subject line."},
            {"name": "body", "required": True, "description": "The email's text."},
        ],
    }],
    "output_format": "json_schema",
    "output_schema": {
        "type": "object",
        "properties": {
            "action": {"type": "string"},
            "due_date": {"type": ["string", "null"]},
        },
        "required": ["action", "due_date"],
        "additionalProperties": False,
    },
    "model_group": {"name": f"{NAME}-models", "models": ["gpt-6-luna", "gpt-6-sol"]},
}

EMAILS = [
    {"subject": "Brook kettle recall: stock count needed",
     "body": "Hi team, the Brook kettle recall goes out next week. Please count the "
             "Brook kettles in the Leeds warehouse and send me the number by "
             "2026-10-21. Thanks, Sam"},
    {"subject": "Thursday's supplier call",
     "body": "Hello, the supplier call on Thursday has moved to the small meeting "
             "room. Could you book the projector for it? Best, Ana"},
]


class Transcript:
    """What the run printed: the page shows `lines`; ids go to `record` only."""

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

    def say(self, line=""):
        print(line)
        self.lines.append(line)


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


def refused(what, r):
    """Stop with the API's own reason. 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')}")


# --- Step 1. Create the task, or reuse it --------------------------------------

def find_task(name):
    """The task's full slug when it already exists and is active, else None.
    The slug is `<account id>.<name>`; the account read gives the prefix. A
    paused task is never reused: a call to it would 404."""
    r = httpx.get(MGMT + "/account", headers=mgmt_headers(), timeout=30)
    if r.status_code != 200:
        refused("account read refused", r)
    slug = r.json()["task_prefix"] + name          # e.g. "3.cookbook-4-action-items"
    r = httpx.get(f"{MGMT}/capabilities/{slug}", headers=mgmt_headers(), timeout=30)
    if r.status_code == 404:
        return None
    if r.status_code != 200:
        refused("task lookup refused", r)
    if not r.json()["is_active"]:
        sys.exit(f"task {slug} exists but is paused; pass a new --name")
    return slug


def create_task(t, name):
    body = dict(TASK, slug=name,
                model_group=dict(TASK["model_group"], name=f"{name[:41]}-models"))
    r = httpx.post(MGMT + "/capabilities/", headers=mgmt_headers(), json=body, timeout=120)
    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(c['model_slug'] for c in task.get('cells') or [])}")
    t.say(f"  stats_pooled: {task.get('stats_pooled')}")
    t.say(f"  exploration.status: {ex.get('status')}")
    fo = ex.get("fans_out") or {}
    t.say(f"  exploration.fans_out: native {fo.get('native')}, /v1 {fo.get('v1')}")
    t.say(f"  exploration.judging: {ex.get('judging')}")
    for line in wrap(ex.get("notice") or "", "  notice: "):
        t.say(line)
    for n in task.get("notices") or []:
        for line in wrap(n, "  notices: "):
            t.say(line)
    return task["slug"]


def wrap(text, first, width=92):
    """Short lines for the page; the words are the API's own."""
    out, line = [], first
    for word in text.split():
        if len(line) + len(word) + 1 > width and line.strip():
            out.append(line.rstrip())
            line = " " * len(first)
        line += word + " "
    if line.strip():
        out.append(line.rstrip())
    return out


# --- Step 2 and 4. The task's progress, from the management API ---------------

def progress(t, slug, when):
    r = httpx.get(f"{MGMT}/capabilities/{slug}", headers=mgmt_headers(), timeout=60)
    if r.status_code != 200:
        refused("task read refused", r)
    m = r.json().get("maturity")
    t.say(f"maturity {when}:")
    if m is None:
        t.say("  (null: not computable right now; read it as unknown, never as matured)")
        return None
    t.say(f"  status {m.get('status')}, {m.get('mature_groups')} of {m.get('total_groups')} "
          f"model groups at {m.get('threshold_band')} or better")
    cd = m.get("countdown_judged_calls")
    t.say(f"  countdown_judged_calls {cd if cd is not None else '(none: no estimate)'}")
    return m


# --- Step 3. Two calls on /v1/chat/completions --------------------------------

HEADERS = ["x-llmbench-model", "x-llmbench-provider", "x-llmbench-band",
           "x-llmbench-phase", "x-llmbench-switch-in", "x-llmbench-cost-usd"]


def call(t, client, slug, email, n):
    """One call in template mode: the task's template renders `prompt_vars`, and
    `messages` stays empty. The x-llmbench-* headers say how the call ran."""
    t.say(f"call {n}: {email['subject']}")
    try:
        raw = client.chat.completions.with_raw_response.create(
            model=slug,
            messages=[],
            extra_body={"llmbench": {"prompt_vars": email}},
        )
    except openai.APIStatusError as e:
        h = e.response.headers
        t.record.setdefault("diagnostics", []).append(h.get("x-llmbench-diagnostic"))
        t.say(f"  failed: HTTP {e.status_code}, {h.get('x-llmbench-error-code')} "
              f"(whose problem: {h.get('x-llmbench-error-class')})")
        return
    except openai.APIConnectionError as e:
        # No response came back. The call may still have run on the server.
        t.say(f"  failed: no response ({type(e).__name__}); the call may still have run")
        return
    h = raw.headers
    answer = json.loads(raw.parse().choices[0].message.content)
    t.say(f"  answer: {json.dumps(answer)}")
    for name in HEADERS:
        t.say(f"  {name + ':':<24} {h.get(name, '(absent)')}")
    t.record.setdefault("request_ids", []).append(h.get("x-llmbench-request-id"))


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--name", default=NAME, help="the task's name (a new name makes a new task)")
    p.add_argument("--page-dir", help="write the transcript here as JSON")
    p.add_argument("--record-json", help="write the task slug and request ids here")
    a = p.parse_args()

    now = dt.datetime.now(dt.timezone.utc)
    t = Transcript()
    t.say(f"run: {now:%Y-%m-%d %H:%M} UTC against {API}")

    slug = find_task(a.name)
    if slug:
        t.say(f"task: {slug} (already there, reused)")
    else:
        slug = create_task(t, a.name)
    t.record["task"] = slug

    progress(t, slug, "before the calls")

    client = OpenAI(base_url=f"{API}/v1", api_key=os.environ["FRONSET_INFERENCE_KEY"],
                    timeout=900, max_retries=0)
    for n, email in enumerate(EMAILS, start=1):
        call(t, client, slug, email, n)

    progress(t, slug, "after the calls")

    if a.page_dir:
        with open(os.path.join(a.page_dir, "cb4_own_task.json"), "w") as f:
            json.dump({"run_date": f"{now:%Y-%m-%d}", "transcript": t.lines}, f, indent=2)
    if a.record_json:
        with open(a.record_json, "w") as f:
            json.dump(t.record, f, indent=2)


if __name__ == "__main__":
    main()