"""The model writes the whole quote: it reads the forms (text) and applies the quoting rules (rules/quote_rules.md),
the rate card and three worked examples (rules/examples.md), all in the prompt (cached). Dev only by default.
Run on your Mac:
    python3 quote_run.py --n=40                  Opus, a fixed stratified sample of 40 dev jobs (~$1)
    python3 quote_run.py                         Opus, all dev jobs (re-uses the 40 already done)
    python3 quote_run.py --model=claude-haiku-5-5
    python3 quote_run.py --photos                forms + site photos (about 1.5x the cost)
    python3 quote_run.py --model=claude-sonnet-5-5 --holdout   the 50 hold-out jobs, once, with the frozen rules
    python3 quote_run.py --model=claude-sonnet-5-5 --final     the 849 unlabelled jobs (the ones that are marked), frozen rules
    python3 quote_run.py --dry                   prompt size and one example prompt, no calls
Resumable. Results: analysis/quote_run_<model>.json. Hold-out and unlabelled jobs are touched only with --holdout / --final.
"""
import base64, json, os, re, subprocess, sys, tempfile, threading, time
from concurrent.futures import ThreadPoolExecutor

ROOT = os.path.dirname(os.path.abspath(__file__))
DRY = "--dry" in sys.argv
KEY = None if DRY else re.search(r"^ANTHROPIC_API_KEY=(\S+)", open(os.path.join(ROOT, ".env")).read(), re.M).group(1)
J = {j["item"]: j for j in json.load(open(os.path.join(ROOT, "jobs.json")))["jobs"]}
SPLIT = json.load(open(os.path.join(ROOT, "splits.json")))
PRICES = {"claude-haiku-5-5": (0.10, 0.50), "claude-sonnet-5-5": (2.0, 10.0), "claude-opus-5-5": (4.0, 20.0)}
MODEL = next((a.split("=", 1)[1] for a in sys.argv if a.startswith("--model=")), "claude-opus-5-5")
PRICE = PRICES[MODEL]; PHOTOS = "--photos" in sys.argv; HOLD = "--holdout" in sys.argv; FINAL = "--final" in sys.argv
TAG = MODEL.split("-")[1] + ("_photos" if PHOTOS else "") + ("_holdout" if HOLD else "") + ("_final" if FINAL else "")
READ_X = 0.05 if MODEL in ("claude-opus-5-5", "claude-sonnet-5-5") else 0.1   # cache read price vs base input (docs: prompt caching)
CACHE = os.path.join(ROOT, "analysis", f"quote_cache_{TAG}.jsonl"); OUT = os.path.join(ROOT, "analysis", f"quote_run_{TAG}.json")
RC = json.load(open(os.path.join(ROOT, "rate_card.json"))); RC = RC["lines"] if isinstance(RC, dict) and "lines" in RC else RC
CODES = {l["code"] for l in RC}
RATE = "\n".join(f"{l['code']} | {l['description']} | per {l.get('unit')} | £{l['unit_price']:g}" for l in RC)

SYSTEM = f"""You are an estimator at a UK drainage and civils contractor. An engineer visited the site and filled in one or more forms. You turn them into a quote using ONLY the rate card below, following the quoting rules. Answer whether it is a manhole job and, if so, what kind (the cover line), and list the quote lines with quantities.
Abbreviations: mh = manhole, ic = inspection chamber, sw = surface water, fw = foul water, ooh = out of hours, aco = channel drain.

RATE CARD (code | description | unit | price):
{RATE}

{open(os.path.join(ROOT, "rules", "quote_rules.md")).read()}

{open(os.path.join(ROOT, "rules", "examples.md")).read()}

{"Site photos follow the forms. Use them for what the forms do not say: the place and what drives over it, the cover style (plate, recessed tray), the surface, the number of items. Several photos often show the same object." if PHOTOS else ""}

Answer with JSON only, in this shape:
{{"job_type":"cover|rebed|gully|channel|other","other_kind":"tanker|pipe_dig|drainage_visit|chamber_repair|plumbing|bespoke_cover|crew_visit|not_applicable","is_manhole":true|false,"kind":"<cover code, or rebed, or null>","lines":[{{"code":"...","qty":1,"rule":"4.3"}}],"colour":"green|yellow|red","check":["what the estimator must confirm"]}}
is_manhole is true for cover and rebed jobs."""


def form_text(job):
    seen, parts = set(), []
    for f in job["remedial_forms"]:
        if not (f.get("issue_type") or f.get("fields")): continue
        k = json.dumps([f.get("issue_type"), f.get("fields")], sort_keys=True)
        if k in seen: continue
        seen.add(k)
        lines = [f"Form {len(parts)+1}: {f.get('issue_type')} / {f.get('request_type')} / {f.get('health_and_safety_rating')}"]
        lines += [f"  {k2}: {v}" for k2, v in (f.get("fields") or {}).items() if v not in (None, "") and k2 != "Type of issue"]
        parts.append("\n".join(lines))
    return "\n".join(parts) or "(no forms)"


lock = threading.Lock(); spent = [0.0]


def call(text):
    body = json.dumps({"model": MODEL, "max_tokens": 5000,
                       "system": [{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}],
                       "messages": [{"role": "user", "content": text}]}).encode()
    for attempt in range(4):
        d = tempfile.mkdtemp(); bf, hf = os.path.join(d, "b.json"), os.path.join(d, "h.txt")
        open(bf, "wb").write(body); open(hf, "w").write(f"x-api-key: {KEY}\nanthropic-version: 2026-10-21\ncontent-type: application/json\n")
        out = subprocess.run(["curl", "-sS", "-m", "180", "https://api.anthropic.com/v1/messages", "-H", f"@{hf}", "--data-binary", f"@{bf}"], capture_output=True, text=True)
        os.remove(hf)
        try: r = json.loads(out.stdout)
        except Exception: time.sleep(5); continue
        if r.get("type") == "error":
            if r["error"].get("type") in ("overloaded_error", "rate_limit_error", "api_error"): time.sleep(10 * (attempt + 1)); continue
            return {"error": r["error"]}, 0
        u = r["usage"]
        cin = u.get("input_tokens", 0) + 1.25 * u.get("cache_creation_input_tokens", 0) + READ_X * u.get("cache_read_input_tokens", 0)
        cost = (cin * PRICE[0] + u["output_tokens"] * PRICE[1]) / 1e6
        txt = "".join(b.get("text", "") for b in r["content"] if b["type"] == "text")
        a, b = txt.find("{"), txt.rfind("}")
        try: ans = json.loads(txt[a:b + 1])
        except Exception: ans = {"raw": txt[:400]}
        if not txt.strip() and attempt < 3: continue
        ans["_unknown_codes"] = [l.get("code") for l in ans.get("lines", []) if l.get("code") not in CODES]
        ans["_tokens"] = [u.get("input_tokens", 0), u.get("cache_creation_input_tokens", 0), u.get("cache_read_input_tokens", 0), u["output_tokens"]]
        return ans, cost
    return {"error": "no valid reply"}, 0


def content(it):
    text = f"FORMS:\n{form_text(J[it])}"
    if not PHOTOS: return f"{text}\n\nWrite the quote. JSON only."
    c = [{"type": "text", "text": text + f"\n\nPHOTOS ({len(J[it]['photos'])}):"}]
    for k, p in enumerate(J[it]["photos"], 1):
        c.append({"type": "text", "text": f"Photo {k}:"})
        c.append({"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": base64.b64encode(open(os.path.join(ROOT, p), "rb").read()).decode()}})
    c.append({"type": "text", "text": "Write the quote. JSON only."})
    return c


def work(it):
    ans, cost = call(content(it))
    rec = {"job": it, "answer": ans, "usd": cost}
    with lock:
        spent[0] += cost
        with open(CACHE, "a") as f: f.write(json.dumps(rec) + "\n")
    return rec


def _ok(rec): return not ("error" in rec["answer"] or "raw" in rec["answer"])
done = {json.loads(l)["job"] for l in open(CACHE) if _ok(json.loads(l))} if os.path.exists(CACHE) else set()   # failed calls are retried
# --n=40: a fixed stratified sample of the dev jobs (proportional by job type, the 3 prompt examples left out), for a cheap first run
N = next((int(a.split("=", 1)[1]) for a in sys.argv if a.startswith("--n=")), None)
pool = [i for i in SPLIT["dev"] if i not in ("J0713", "J0705", "J0243")]
if HOLD:                                   # the 50 hold-out jobs, run once with the frozen rules
    pool = list(SPLIT["holdout"])
if FINAL:                                  # all jobs without a label (999 - 150), the ones Metcor marks
    pool = sorted(i for i in J if i not in set(SPLIT["dev"]) | set(SPLIT["holdout"]))
if N:
    sf = os.path.join(ROOT, "analysis", f"s5_sample{N}.json")
    if os.path.exists(sf):
        pool = json.load(open(sf))
    else:
        import random
        from collections import defaultdict
        LABT = {e["item"]: e["job_type"].split(":")[0] + ":" + e["job_type"].split(":")[1].split(",")[0][:25] for e in json.load(open(os.path.join(ROOT, "labelled_examples.json")))["jobs"]}
        grp = defaultdict(list)
        for i in pool: grp[LABT[i]].append(i)
        rnd = random.Random(7); pick = []
        for g, ids in sorted(grp.items()):
            rnd.shuffle(ids); pick += ids[:max(1, round(N * len(ids) / len(pool)))]
        pool = sorted(pick[:N] if len(pick) >= N else pick)
        json.dump(pool, open(sf, "w"))
todo = [i for i in pool if i not in done]
if DRY:
    print(f"system prompt: {len(SYSTEM)} characters (~{len(SYSTEM)//4} tokens)\n\nFORMS:\n{form_text(J[todo[0]])}\n\n{len(todo)} jobs to do"); sys.exit()
print(f"model {MODEL}, rules + rate card + examples in the prompt: {len(todo)} jobs to do ({len(done)} already done)", flush=True)
first = todo[:1]; rest = todo[1:]
t0 = time.time()
for it in first: work(it)                                  # writes the prompt cache once, then the rest read it
with ThreadPoolExecutor(8) as ex:
    for n, rec in enumerate(ex.map(work, rest), 2):
        if n % 10 == 0 or n == len(todo): print(f"  {n}/{len(todo)} done, ${spent[0]:.2f}, {time.time()-t0:.0f}s", flush=True)
recs = {}
for l in open(CACHE):                       # keep a good answer over a failed one
    r = json.loads(l)
    if r["job"] not in recs or _ok(r): recs[r["job"]] = r
json.dump({"model": MODEL, "records": recs, "total_usd": sum(r["usd"] for r in recs.values())}, open(OUT, "w"), indent=1)
bad = sum(1 for r in recs.values() if "error" in r["answer"] or "raw" in r["answer"])
unk = sum(len(r["answer"].get("_unknown_codes", [])) for r in recs.values())
T = [r["answer"].get("_tokens") for r in recs.values() if r["answer"].get("_tokens")]
print(f"\nPrompt cache: written on {sum(t[1] > 0 for t in T)} calls, read on {sum(t[2] > 0 for t in T)} of {len(T)} calls"
      + (" -> NOT cached (both fields 0)" if not any(t[1] or t[2] for t in T) else ""))
print(f"Saved {len(recs)} answers to {os.path.relpath(OUT, ROOT)}. Unusable: {bad}. Codes not on the rate card: {unk}. Total cost: ${sum(r['usd'] for r in recs.values()):.2f}")