"""Thin Claude API wrapper: structured outputs, provenance ledger, cost accounting, budget cap, response cache."""
from __future__ import annotations
import json, time, hashlib, threading
from pathlib import Path
import anthropic
from . import config
from .provenance import Ledger, sha256_text, now_iso, code_version, ROOT

_lock = threading.Lock()
LEDGER = Ledger(ROOT / "data" / "annotations" / "api_ledger.jsonl")
CACHE_DIR = ROOT / "data" / "annotations" / "_llm_cache"; CACHE_DIR.mkdir(parents=True, exist_ok=True)

def spent_usd() -> float:
    return round(sum(r.get("cost_usd", 0.0) for r in LEDGER.read()), 4)

def cost_of(model: str, usage) -> float:
    pr = config.load("models")["pricing"].get(model, {"input": 4.0, "output": 20.0, "cache_read": 0.2, "cache_write": 5.0})
    cr = getattr(usage, "cache_read_input_tokens", 0) or 0; cw = getattr(usage, "cache_creation_input_tokens", 0) or 0
    return (usage.input_tokens * pr["input"] + usage.output_tokens * pr["output"] + cr * pr["cache_read"] + cw * pr["cache_write"]) / 1e6

class BudgetExceeded(RuntimeError): pass

def extract_json(text: str):
    """Parse a JSON object from model text (tolerates code fences / surrounding prose)."""
    t = text.strip()
    if t.startswith("```"): t = t.split("\n", 1)[1] if "\n" in t else t; t = t.rsplit("```", 1)[0]
    try: return json.loads(t)
    except json.JSONDecodeError: pass
    i, j = t.find("{"), t.rfind("}")
    if i >= 0 and j > i:
        try: return json.loads(t[i:j + 1])
        except json.JSONDecodeError: return None
    return None

def call(system: str, user: str, *, model: str | None = None, schema: dict | None = None, max_tokens: int = 8000, effort: str = "medium",
         purpose: str = "", prompt_id: str = "", prompt_version: str = "", use_cache: bool = True, meta: dict | None = None, parse_json: bool = False) -> dict:
    """One Messages API call. Returns {"text", "json", "usage", "cost_usd", "cached", "request_id", "model"}; logs to the ledger."""
    cfg = config.load("models"); model = model or cfg["annotation_model"]
    key = sha256_text(json.dumps({"m": model, "s": system, "u": user, "schema": schema, "eff": effort, "mt": max_tokens}, sort_keys=True))
    cp = CACHE_DIR / f"{key}.json"
    if use_cache and cp.exists():
        d = json.loads(cp.read_text()); d["cached"] = True; return d
    with _lock:
        if spent_usd() >= cfg["budget_usd_cap"]: raise BudgetExceeded(f"budget cap {cfg['budget_usd_cap']} USD reached (spent {spent_usd()})")
    client = anthropic.Anthropic(api_key=config.api_key(), max_retries=3, timeout=600)
    kw = dict(model=model, max_tokens=max_tokens, system=[{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}],
              messages=[{"role": "user", "content": user}], output_config={"effort": effort})
    if schema: kw["output_config"]["format"] = {"type": "json_schema", "schema": schema}
    t0 = time.time(); last = None
    for attempt in range(3):
        try:
            r = client.messages.create(**kw); break
        except (anthropic.RateLimitError, anthropic.APIStatusError, anthropic.APIConnectionError) as e:
            last = e; time.sleep(5 * (attempt + 1))
    else:
        raise last
    text = "".join(b.text for b in r.content if b.type == "text")
    out = {"text": text, "json": None, "model": r.model, "stop_reason": r.stop_reason, "latency_s": round(time.time() - t0, 2),
           "usage": {"input": r.usage.input_tokens, "output": r.usage.output_tokens, "cache_read": getattr(r.usage, "cache_read_input_tokens", 0) or 0, "cache_write": getattr(r.usage, "cache_creation_input_tokens", 0) or 0},
           "cost_usd": round(cost_of(model, r.usage), 6), "request_id": getattr(r, "_request_id", None), "cached": False, "ts": now_iso(), "effort": effort}
    if schema or parse_json:
        out["json"] = extract_json(text)
    if r.stop_reason == "refusal": out["refusal"] = getattr(r, "stop_details", None) and r.stop_details.category
    LEDGER.append({"purpose": purpose, "prompt_id": prompt_id, "prompt_version": prompt_version, "model": r.model, "effort": effort, "usage": out["usage"], "cost_usd": out["cost_usd"],
                   "latency_s": out["latency_s"], "request_id": out["request_id"], "stop_reason": r.stop_reason, "cache_key": key, "code_version": code_version(), **(meta or {})})
    if use_cache: cp.write_text(json.dumps(out, ensure_ascii=False))
    return out