"""Anthropic Messages API wrapper for WP20 extraction.

Two modes share one request builder:
  - direct(): synchronous call (smoke tests, retries, ≤ 5 tries), cached by request hash.
  - batch(): Message Batches API (50 % price). Requests are written to results/raw/batches/<name>.requests.jsonl
    and results to results/raw/batches/<name>.results.jsonl; nothing is lost if the process dies.
Secrets: ANTHROPIC_API_KEY from the environment only (loaded from a .env OUTSIDE the repo if WP20_ENV_FILE is set).
Usage accounting: every response's usage block is appended to results/raw/api_usage.jsonl.
"""
from __future__ import annotations

import hashlib
import json
import os
import time
from pathlib import Path

import anthropic
import yaml

ROOT = Path(__file__).resolve().parents[2]
RAW = ROOT / "results" / "raw"
RAW.mkdir(parents=True, exist_ok=True)
(RAW / "batches").mkdir(exist_ok=True)
CACHE = ROOT / "data" / "cache"
CACHE.mkdir(parents=True, exist_ok=True)

PROMPTS = yaml.safe_load(open(ROOT / "configs" / "prompts.yaml"))
MODELS = yaml.safe_load(open(ROOT / "configs" / "models.yaml"))


def _load_env() -> None:
    f = os.environ.get("WP20_ENV_FILE")
    if f and os.path.exists(f) and "ANTHROPIC_API_KEY" not in os.environ:
        for line in open(f, encoding="utf-8"):
            if line.startswith("ANTHROPIC_API_KEY="):
                os.environ["ANTHROPIC_API_KEY"] = line.split("=", 1)[1].strip()


_load_env()
client = anthropic.Anthropic(max_retries=3, timeout=180)


def model_params(tier: str) -> dict:
    """Per-model request parameters (thinking / effort) from configs/models.yaml."""
    m = MODELS["tiers"][tier]
    p = {"model": m["id"], "max_tokens": m["max_tokens"]}
    if m.get("thinking"):
        p["thinking"] = m["thinking"]
    if m.get("effort"):
        p["output_config"] = {"effort": m["effort"]}
    return p


def image_request(tier: str, url: str, schema: dict) -> dict:
    p = model_params(tier)
    p["system"] = PROMPTS["image_system"]
    p["messages"] = [{"role": "user", "content": [
        {"type": "image", "source": {"type": "url", "url": url}},
        {"type": "text", "text": PROMPTS["image_user"]}]}]
    oc = p.get("output_config", {})
    oc["format"] = {"type": "json_schema", "schema": schema}
    p["output_config"] = oc
    return p


def text_request(tier: str, text: str, schema: dict) -> dict:
    p = model_params(tier)
    p["system"] = PROMPTS["text_system"]
    p["messages"] = [{"role": "user", "content": PROMPTS["text_user"].format(text=text)}]
    oc = p.get("output_config", {})
    oc["format"] = {"type": "json_schema", "schema": schema}
    p["output_config"] = oc
    return p


def _hash(params: dict) -> str:
    return hashlib.sha256(json.dumps(params, sort_keys=True, ensure_ascii=False).encode()).hexdigest()[:24]


def _log_usage(tag: str, model: str, usage, extra: dict | None = None) -> None:
    rec = {"ts": time.time(), "tag": tag, "model": model,
           "input_tokens": getattr(usage, "input_tokens", None), "output_tokens": getattr(usage, "output_tokens", None),
           "cache_read": getattr(usage, "cache_read_input_tokens", None)}
    if extra:
        rec.update(extra)
    with open(RAW / "api_usage.jsonl", "a") as fh:
        fh.write(json.dumps(rec) + "\n")


def parse_json_text(message) -> dict | None:
    if getattr(message, "stop_reason", None) == "refusal":
        return None
    for b in message.content:
        if b.type == "text":
            try:
                return json.loads(b.text)
            except json.JSONDecodeError:
                return None
    return None


def direct(params: dict, tag: str = "direct") -> dict:
    """Synchronous call with on-disk cache. Returns {'ok', 'data', 'usage', 'model', 'stop_reason', 'error'}."""
    h = _hash(params)
    cf = CACHE / f"{h}.json"
    if cf.exists():
        return json.load(open(cf))
    try:
        msg = client.messages.create(**params)
    except anthropic.APIStatusError as e:
        return {"ok": False, "error": f"{e.status_code}: {e.message}", "hash": h}
    data = parse_json_text(msg)
    out = {"ok": data is not None, "data": data, "model": msg.model, "stop_reason": msg.stop_reason,
           "usage": {"input_tokens": msg.usage.input_tokens, "output_tokens": msg.usage.output_tokens}, "hash": h}
    _log_usage(tag, msg.model, msg.usage)
    json.dump(out, open(cf, "w"), ensure_ascii=False)
    return out


# ---------------------------------------------------------------- batches
def batch_submit(name: str, requests: list[tuple[str, dict]]) -> list[str]:
    """Submit requests [(custom_id, params)] in chunks (≤ 10,000 per batch, ≤ 200 MB). Returns batch ids."""
    from anthropic.types.message_create_params import MessageCreateParamsNonStreaming
    from anthropic.types.messages.batch_create_params import Request

    reqf = RAW / "batches" / f"{name}.requests.jsonl"
    with open(reqf, "w") as fh:
        for cid, p in requests:
            fh.write(json.dumps({"custom_id": cid, "params": p}, ensure_ascii=False) + "\n")
    ids = []
    chunk = 4000
    for i in range(0, len(requests), chunk):
        part = requests[i:i + chunk]
        b = client.messages.batches.create(requests=[
            Request(custom_id=cid, params=MessageCreateParamsNonStreaming(**p)) for cid, p in part])
        ids.append(b.id)
        with open(RAW / "batches" / f"{name}.batch_ids.txt", "a") as fh:
            fh.write(f"{b.id}\t{len(part)}\t{time.strftime('%Y-%m-%dT%H:%M:%S')}\n")
    return ids


def batch_status(batch_id: str):
    return client.messages.batches.retrieve(batch_id)


def batch_collect(name: str, batch_id: str) -> int:
    """Append results of an ended batch to <name>.results.jsonl (idempotent by custom_id). Returns n written."""
    resf = RAW / "batches" / f"{name}.results.jsonl"
    seen = set()
    if resf.exists():
        for line in open(resf):
            try:
                seen.add(json.loads(line)["custom_id"])
            except Exception:
                pass
    n = 0
    with open(resf, "a") as fh:
        for r in client.messages.batches.results(batch_id):
            if r.custom_id in seen:
                continue
            rec = {"custom_id": r.custom_id, "type": r.result.type, "batch_id": batch_id}
            if r.result.type == "succeeded":
                msg = r.result.message
                rec["model"] = msg.model
                rec["stop_reason"] = msg.stop_reason
                rec["data"] = parse_json_text(msg)
                rec["usage"] = {"input_tokens": msg.usage.input_tokens, "output_tokens": msg.usage.output_tokens}
                _log_usage(f"batch:{name}", msg.model, msg.usage)
            else:
                err = getattr(r.result, "error", None)
                inner = getattr(err, "error", None) if err else None
                rec["error"] = getattr(inner, "type", None) or (getattr(err, "type", None) if err else None)
                rec["error_message"] = getattr(inner, "message", None) if inner else (str(err) if err else None)
            fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
            n += 1
    return n