#!/usr/bin/env python3
"""Fronset cookbook 3: a batch of 50 lines on a settled catalogue task.

Summarise The Federalist Nos. 1-50 in one batch. The task has settled (we know
which model clears its bar), so every line goes to that one model on the
provider's batch lane. The script uploads the file, creates the batch, polls it,
downloads the output and error files, and reads each line's model, service tier
and cost. The fee arithmetic for larger volumes is on the cookbook page.

    pip install openai          # httpx comes with it
    export FRONSET_INFERENCE_KEY=llmb_i_...     # /v1: files and batches
    export FRONSET_MANAGEMENT_KEY=llmb_m_...    # task read, rate quote, usage
    python cb3_batch_settled_task.py --out-dir ./cb3_out

Keys are read from the environment only. Nothing here prints a key.

The inputs are The Federalist Nos. 1-50 (1787-88, public domain), from Project
Gutenberg eBook #1404.
"""
import argparse
import collections
import datetime as dt
import hashlib
import hmac
import io
import json
import os
import re
import sys
import time

import httpx
import openai
from openai import OpenAI

API = os.environ.get("FRONSET_API_BASE", "https://api.fronset.ai")
TASK = "structured-content-summarization"
SOURCE_URL = "https://www.gutenberg.org/cache/epub/1404/pg1404.txt"
SOURCE_PAGE = "https://www.gutenberg.org/ebooks/1404"

# The per-call fee card (Pricing Terms §3): 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)]

# Terminal batch statuses: polling stops at any of these.
TERMINAL = {"completed", "failed", "expired", "cancelled"}


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 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 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. The inputs: one line per paper, in template mode ---------------

def federalist_papers(count):
    """[(number, subtitle, text)] for The Federalist Nos. 1..count."""
    raw = httpx.get(SOURCE_URL, timeout=60, follow_redirects=True).content
    text = raw.decode("utf-8-sig").replace("\r\n", "\n")
    start = text.find("*** START OF")
    end = text.find("*** END OF")
    text = text[text.find("\n", start) + 1:end] if start >= 0 and end > start else text
    parts = re.split(r"(?m)^FEDERALIST No\. (\d+)\s*$", text)
    papers, seen = [], set()
    for number, body in zip(parts[1::2], parts[2::2]):
        n = int(number)
        if n in seen or n > count:
            continue
        seen.add(n)
        lines = [ln.strip() for ln in body.strip().split("\n")]
        subtitle = next((ln for ln in lines if ln), "")
        marker = "To the People of the State of New York:"
        essay = body.split(marker, 1)[1] if marker in body else body
        papers.append((n, subtitle, essay.strip()))
    papers.sort()
    return papers, hashlib.sha256(raw).hexdigest()


def batch_file(papers):
    """One JSONL line per paper. `model` is not on the lines: it is hoisted on
    the create call, so it cannot disagree between lines."""
    out = io.BytesIO()
    for n, subtitle, essay in papers:
        line = {"custom_id": f"federalist-{n:02d}", "method": "POST",
                "url": "/v1/chat/completions",
                "body": {"messages": [],
                         "llmbench": {"prompt_vars": {
                             "content_text": essay,
                             "content_title": f"The Federalist No. {n}: {subtitle}",
                             "content_type": "essay",
                             "content_url": SOURCE_PAGE,
                             "summary_structure": "## Summary\nTwo sentences: the "
                                                  "paper's claim, and its main argument.",
                         }}}}
        out.write((json.dumps(line) + "\n").encode("utf-8"))
    return out.getvalue()


# --- Step 5 (optional). Verifying a webhook on your own server -------------

def verify_webhook(secret, timestamp, raw_body, signature_header):
    """True when an X-Signature header is ours. HMAC-SHA256 over the RAW bytes
    "<X-Timestamp>." + body, with the endpoint's secret. Verify before parsing.
    This script polls instead; this is the check your webhook receiver runs."""
    expected = hmac.new(secret.encode(), f"{timestamp}.".encode() + raw_body,
                        hashlib.sha256).hexdigest()
    return hmac.compare_digest(f"sha256={expected}", signature_header or "")


def task_state(t):
    """The task's maturity and its batch flag, from the management API."""
    r = httpx.get(f"{API}/api/inference/v1/management/capabilities/{TASK}",
                  headers=mgmt_headers(), timeout=60)
    if r.status_code != 200:
        refused("task read refused", r)
    task = r.json()
    m = task.get("maturity") or {}
    t.say(f"task: {TASK}")
    t.say(f"  maturity      {m.get('status', '(absent)')} "
          f"({m.get('mature_groups')} of {m.get('total_groups')} model groups at "
          f"{m.get('threshold_band')} or better)")
    t.say(f"  batch_eligible {task.get('batch_eligible')}")
    return task


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 batch_rate(task):
    """The published batch-lane rate a line of this task resolves to."""
    mg = f"{API}/api/inference/v1/management"
    params = {"lane": "batch", "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/{TASK}/quote", headers=mgmt_headers(),
                  timeout=30, params=params)
    if r.status_code == 404:
        return None
    if r.status_code != 200:
        refused("rate quote refused", r)
    q = r.json()
    res = q.get("resolved") or {}
    if not (q.get("priced") and res.get("input") and res.get("output")):
        return None
    return res


def read_lines(client, file_id, keep_as):
    """Download one result file, keep a local copy, and parse it."""
    if not file_id:
        return []
    raw = client.files.content(file_id).text
    if keep_as:
        with open(keep_as, "w") as f:
            f.write(raw)
    return [json.loads(ln) for ln in raw.splitlines() if ln.strip()]


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--count", type=int, default=50, help="papers to summarise (1-85)")
    p.add_argument("--poll-seconds", type=int, default=30)
    p.add_argument("--out-dir", help="keep the input, output and error files here")
    p.add_argument("--page-dir", help="write the transcript here as JSON")
    p.add_argument("--record-json", help="write the file and batch ids here")
    a = p.parse_args()
    today = dt.datetime.now(dt.timezone.utc).date()
    t = Transcript()
    t.say(f"run: {dt.datetime.now(dt.timezone.utc):%Y-%m-%d %H:%M} UTC against {API}/v1")

    # No automatic retries: a retried create would make a second batch, and
    # every line of it would run and be billed again.
    client = OpenAI(base_url=f"{API}/v1", api_key=os.environ["FRONSET_INFERENCE_KEY"],
                    timeout=120, max_retries=0)
    task = task_state(t)

    # Step 1. The input file.
    papers, digest = federalist_papers(a.count)
    data = batch_file(papers)
    t.say()
    t.say(f"inputs: The Federalist Nos. {papers[0][0]}-{papers[-1][0]} ({len(papers)} lines, "
          f"{len(data):,} bytes; source sha256 {digest[:12]})")
    if a.out_dir:
        os.makedirs(a.out_dir, exist_ok=True)
        with open(os.path.join(a.out_dir, "cb3_input.jsonl"), "wb") as f:
            f.write(data)

    # Step 2. Upload it.
    upload = client.files.create(file=("cb3_input.jsonl", data), purpose="batch")
    t.record["input_file_id"] = upload.id

    # Step 3. Create the batch. `model` is hoisted into the llmbench block.
    before = calls_this_month(today)
    batch = client.batches.create(
        input_file_id=upload.id,
        endpoint="/v1/chat/completions",
        completion_window="24h",
        metadata={"cookbook": "3"},
        extra_body={"llmbench": {"model": TASK}},
    )
    t.record["batch_id"] = batch.id
    created = time.monotonic()
    t.say(f"batch created: status {batch.status}, expires "
          f"{dt.datetime.fromtimestamp(batch.expires_at, dt.timezone.utc):%Y-%m-%d %H:%M} UTC")

    # Step 4. Poll. Polling is authoritative; a webhook is a nudge to poll now.
    t.say()
    t.say("polling")
    last = None
    while True:
        try:
            batch = client.batches.retrieve(batch.id)
        except (openai.APIConnectionError, openai.InternalServerError) as e:
            # A read is safe to repeat: wait and poll again.
            print(f"  poll failed ({type(e).__name__}); trying again")
            time.sleep(a.poll_seconds)
            continue
        if batch.status != last:
            rc = batch.request_counts
            t.say(f"  +{int(time.monotonic() - created):>5}s  {batch.status:<11} "
                  f"{rc.completed if rc else 0}/{rc.total if rc else 0} done, "
                  f"{rc.failed if rc else 0} failed")
            last = batch.status
        if batch.status in TERMINAL:
            break
        time.sleep(a.poll_seconds)
    elapsed = int(time.monotonic() - created)
    extra = (batch.model_extra or {}).get("llmbench") or {}
    t.say(f"finished in {elapsed // 60} min {elapsed % 60} s; "
          f"executed lane: {extra.get('executed_lane', '(not reported)')}")

    # Step 6. Download both files and read every line.
    out_dir = a.out_dir
    outputs = read_lines(client, batch.output_file_id,
                         out_dir and os.path.join(out_dir, "cb3_output.jsonl"))
    errors = read_lines(client, batch.error_file_id,
                        out_dir and os.path.join(out_dir, "cb3_error.jsonl"))
    t.record.update(output_file_id=batch.output_file_id, error_file_id=batch.error_file_id)

    models, tiers = collections.Counter(), collections.Counter()
    list_cost, tokens_in, tokens_out, sample = 0.0, 0, 0, None
    for line in outputs:
        body = (line.get("response") or {}).get("body") or {}
        mirror = body.get("llmbench") or {}
        models[f"{mirror.get('model')} ({mirror.get('provider')})"] += 1
        tiers[mirror.get("service-tier") or "(absent)"] += 1
        list_cost += float(mirror.get("cost-usd") or 0)
        usage = body.get("usage") or {}
        tokens_in += usage.get("prompt_tokens") or 0
        tokens_out += usage.get("completion_tokens") or 0
        if sample is None and body.get("choices"):
            sample = (line["custom_id"], body["choices"][0]["message"]["content"])
    error_codes = collections.Counter()
    for line in errors:
        err = line.get("error") or (((line.get("response") or {}).get("body")) or {}).get("error") or {}
        error_codes[err.get("code") or "(no code)"] += 1

    t.say()
    t.say(f"lines: {len(outputs)} in the output file, {len(errors)} in the error file")
    for code, n in error_codes.most_common():
        t.say(f"  error {code}: {n}")
    t.say("answered by")
    for name, n in models.most_common():
        t.say(f"  {name:<44} {n} of {len(outputs)}")
    t.say("service tier")
    for name, n in tiers.most_common():
        t.say(f"  {name:<44} {n} of {len(outputs)}")
    if sample:
        text = sample[1].strip().replace("\n", " ")
        t.say(f"first answer ({sample[0]}): {text[:240]}{'...' if len(text) > 240 else ''}")

    # Step 7. What it cost.
    t.say()
    t.say("what this run cost us")
    t.say(f"  tokens               {tokens_in:,} in, {tokens_out:,} out")
    t.say(f"  list cost (sum of cost-usd)  ${list_cost:.6f}  (the answering model's "
          "list cost per line; not the bill)")
    rate = batch_rate(task)
    if rate:
        r_in, r_out = float(rate["input"]["rate"]), float(rate["output"]["rate"])
        t.say(f"  published batch rate ${r_in:.3f} in / ${r_out:.3f} out per 1M tokens "
              f"(quality {rate['quality_step']}, floor {rate['eligibility_floor']}, since "
              f"{rate['output']['effective_from'][:10]})")
        t.say(f"  model usage          ${(tokens_in * r_in + tokens_out * r_out) / 1e6:.6f}  "
              "(the tokens at the published batch rate)")
    else:
        t.say("  published batch rate NONE in force for this shape: billed at cost")
    done = len(outputs)
    t.say(f"  per-call fee         ${per_call_fee(before, done):.4f}  ({done} calls, numbers "
          f"{before + 1:,}-{before + done:,} 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")

    if a.page_dir:
        with open(os.path.join(a.page_dir, "cb3_batch.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(t.record, f, indent=2)


if __name__ == "__main__":
    main()