#!/usr/bin/env python3
"""
Does an organization-level "keep answers short" instruction actually save tokens?

A controlled three-arm test.

  Arm A  no system prompt at all                          (baseline)
  Arm B  a full org instruction block, LENGTH line REMOVED (isolates the block's own cost)
  Arm C  the same block, LENGTH line present               (the deployed state)

B vs C isolates the single instruction line. A vs C gives the honest end-to-end
effect, including the input tokens the block itself adds to every request.

20 fixed prompts x 3 arms x 3 repetitions = 180 calls. One CSV row per call.

The org block below is a generic, publishable stand-in modelled on a real one
deployed at a manufacturing client. Same structure, same section set, same order
of magnitude in length. Meridian Industrial is invented. No client content.

Usage:  python3 token_measure.py [--reps N] [--out results.csv]
Key:    ~/Documents/GitHub/anthropic_credentials.json  ->  {"api_key": "..."}
"""

import csv
import json
import os
import subprocess
import sys
import time

MODEL = "claude-sonnet-5"
MAX_TOKENS = 2048
API = "https://api.anthropic.com/v1/messages"
KEY_FILE = os.path.expanduser("~/Documents/GitHub/anthropic_credentials.json")

# ---------------------------------------------------------------------------
# The single line under test. Present in arm C, absent in arm B.
# ---------------------------------------------------------------------------
LENGTH_LINE = (
    "LENGTH: match the reply to the question. Short question, short answer. "
    "No preamble, no recap. Go longer when the work needs it or you're asked."
)

# ---------------------------------------------------------------------------
# The org block, minus the LENGTH line. Arm B uses this as-is.
# Arm C inserts LENGTH_LINE at the marker.
# ---------------------------------------------------------------------------
BLOCK_TEMPLATE = """You are Claude, working inside Meridian Industrial. Serve the One Meridian mission: exceptional service and value for employees, customers, and communities.

WHO YOU'RE HELPING, every session: if you can tell who the person is, greet them by name and role and confirm. If not, ask. Sign work as the actual person, never as anyone else.

MERIDIAN IN ONE LINE: a B2B industrial components manufacturer, about 1,300 people across the US and Canada, private-equity backed. Product range: meridian-industrial.example/products/. Transformation plan: Project LIFTOFF.

THE ONE MERIDIAN BEHAVIORS, reinforce and flag anything against them: Put People First; Train and Develop; Execute Consistently; Set the Right Customer Expectation; Accountable to Each Other; Act on Insights; Leaders Are Here to Serve.

HOW TO WRITE: lead with the answer. Short, direct sentences. No filler openers. No em-dashes. Avoid: synergy, leverage, ecosystem, circle back. Internal tone concise; external polished; sign external drafts "Kind regards." Do not write like AI.
{LENGTH_SLOT}
NON-NEGOTIABLES:
- Do not fabricate. If you do not know a fact, number, name, or file path, say so and ask.
- Never move, rename, delete, or change the extension of any Meridian file or folder. Thousands of formulas depend on the paths, and Old/Archive folders are load-bearing. If asked to "clean up", decline and explain.
- Every output is a draft. A person reviews and approves it before it is sent, filed or acted on.
- Keep restricted data (PII, PHI, privileged, financials, HR records) out of any tool not approved for it. You have only the signed-in person's access; do not work around it.
- Never produce financial statements, audit documents, regulatory filings, or employment decisions without documented human review. Report suspected violations or data exposure to the IT Service Desk straight away.

BRANDED WORK: any document, deck or chart a Meridian person will present or share must follow the meridian-branding skill. Never invent colors, fonts, logos or layout.

SKILLS AND DEPTH: the company skills, the department files and the AI Governance Policy are authoritative. All live in the company SharePoint library as CLAUDE.md, skills/, reference/branding/, Departments/, AI Documentation/. Load a department file only when the work concerns it. If a skill is not already loaded, read it from the library via the Microsoft 365 connector first. If the connector is off, say so and ask the person to connect it rather than guessing.

SETUP CHECKS: if asked for the Meridian AI check phrase, answer exactly: cobalt-falcon-47.

DEFAULTS: cite a source for factual claims. Put missing data at the top of a deliverable."""

ARM_B_SYSTEM = BLOCK_TEMPLATE.replace("{LENGTH_SLOT}", "")
ARM_C_SYSTEM = BLOCK_TEMPLATE.replace("{LENGTH_SLOT}", "\n" + LENGTH_LINE + "\n")

ARMS = {
    "A": None,
    "B": ARM_B_SYSTEM,
    "C": ARM_C_SYSTEM,
}

# ---------------------------------------------------------------------------
# 20 fixed prompts. Ordinary business questions, the kind a non-technical
# employee actually types. Published with the results so the test is
# reproducible rather than cherry-picked.
# ---------------------------------------------------------------------------
PROMPTS = [
    ("factual",    "What is the difference between gross margin and contribution margin?"),
    ("factual",    "What does DSO mean in finance and how is it calculated?"),
    ("howto",      "How do I create a pivot table in Excel from a sales export?"),
    ("howto",      "How do I set up an out of office reply in Outlook?"),
    ("draft",      "Draft a short email to a supplier asking them to confirm a delivery date that has slipped by a week."),
    ("draft",      "Write a two sentence note to my team letting them know Friday's meeting is moving to Monday."),
    ("summary",    "Summarise the key risks a manufacturer faces when a single customer is more than 20 percent of revenue."),
    ("summary",    "Give me the main points of a standard net 30 payment term and what it means for cash flow."),
    ("compare",    "What is the difference between a purchase order and an invoice?"),
    ("compare",    "Should a small team use a shared mailbox or a distribution list? Compare them."),
    ("policy",     "Can I put customer pricing data into a public AI chatbot?"),
    ("policy",     "What should I do if I think I sent a spreadsheet to the wrong external address?"),
    ("rewrite",    "Rewrite this to be clearer: 'Per our previous correspondence regarding the aforementioned matter, we would like to circle back and touch base at your earliest convenience.'"),
    ("rewrite",    "Make this shorter without losing meaning: 'It has come to our attention that there may potentially be some issues with the current process which we believe could be improved upon.'"),
    ("analysis",   "Our on time delivery dropped from 94 percent to 88 percent in one quarter. What are the most likely causes to check first?"),
    ("analysis",   "What questions should I ask before approving a capital expenditure request for new equipment?"),
    ("shortfact",  "What is EBITDA?"),
    ("shortfact",  "What day of the week does a calendar quarter usually end on?"),
    ("openended",  "How should I structure a weekly one to one with someone who reports to me?"),
    ("openended",  "What makes a good internal process document?"),
]


def load_key():
    with open(KEY_FILE) as f:
        return json.load(f)["api_key"]


def call(api_key, system, user_text):
    """One Messages API call via curl. Returns (usage_dict, text, latency_s)."""
    payload = {
        "model": MODEL,
        "max_tokens": MAX_TOKENS,
        "messages": [{"role": "user", "content": user_text}],
    }
    if system:
        payload["system"] = system

    cmd = [
        "curl", "-sS", "-X", "POST", API,
        "-H", "content-type: application/json",
        "-H", "anthropic-version: 2023-06-01",
        "-H", f"x-api-key: {api_key}",
        "--data-binary", "@-",
    ]
    t0 = time.time()
    proc = subprocess.run(cmd, input=json.dumps(payload), capture_output=True, text=True)
    latency = time.time() - t0

    if proc.returncode != 0:
        raise RuntimeError(f"curl failed: {proc.stderr[:300]}")
    try:
        data = json.loads(proc.stdout)
    except json.JSONDecodeError:
        raise RuntimeError(f"bad JSON: {proc.stdout[:300]}")
    if "error" in data:
        raise RuntimeError(f"API error: {data['error']}")

    text = "".join(b.get("text", "") for b in data.get("content", []))
    return data.get("usage", {}), text, latency


def main():
    reps = 3
    out_path = "token-measure-results.csv"
    args = sys.argv[1:]
    if "--reps" in args:
        reps = int(args[args.index("--reps") + 1])
    if "--out" in args:
        out_path = args[args.index("--out") + 1]

    api_key = load_key()

    print(f"model            {MODEL}")
    print(f"max_tokens       {MAX_TOKENS}")
    print(f"arm B system     {len(ARM_B_SYSTEM):,} chars")
    print(f"arm C system     {len(ARM_C_SYSTEM):,} chars  (+{len(ARM_C_SYSTEM)-len(ARM_B_SYSTEM)} for the LENGTH line)")
    print(f"prompts          {len(PROMPTS)}")
    print(f"reps             {reps}")
    print(f"total calls      {len(PROMPTS) * len(ARMS) * reps}")
    print(f"out              {out_path}")
    print()

    fields = [
        "arm", "rep", "prompt_id", "category", "prompt",
        "input_tokens", "output_tokens",
        "cache_creation_input_tokens", "cache_read_input_tokens",
        "latency_s", "response",
    ]
    n = 0
    total = len(PROMPTS) * len(ARMS) * reps
    failures = 0

    with open(out_path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for rep in range(1, reps + 1):
            for pid, (cat, prompt) in enumerate(PROMPTS):
                for arm, system in ARMS.items():
                    n += 1
                    try:
                        usage, text, latency = call(api_key, system, prompt)
                    except Exception as e:
                        failures += 1
                        print(f"[{n}/{total}] arm {arm} p{pid} FAILED: {e}")
                        continue
                    w.writerow({
                        "arm": arm,
                        "rep": rep,
                        "prompt_id": pid,
                        "category": cat,
                        "prompt": prompt,
                        "input_tokens": usage.get("input_tokens", ""),
                        "output_tokens": usage.get("output_tokens", ""),
                        "cache_creation_input_tokens": usage.get("cache_creation_input_tokens", 0),
                        "cache_read_input_tokens": usage.get("cache_read_input_tokens", 0),
                        "latency_s": round(latency, 2),
                        "response": text,
                    })
                    f.flush()
                    print(f"[{n}/{total}] arm {arm} rep {rep} p{pid:<2} "
                          f"in={usage.get('input_tokens'):<5} out={usage.get('output_tokens'):<5} "
                          f"{latency:.1f}s  {cat}")

    print()
    print(f"done. {n - failures} rows written, {failures} failures -> {out_path}")


if __name__ == "__main__":
    main()
