"""Stream-json reader and accounting for one `claude -p` call: pure functions, no IO."""
import datetime as _dt
import json
import re
import time

TTL_1H, TTL_5M = 3600, 300
CACHE_HIT_RATIO = 0.9
MU_FIELDS = ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")
TOKEN_KEYS = ("input", "output", "cache_creation", "cache_read", "cache_creation_1h", "cache_creation_5m")


def iso(t: float | None = None) -> str:
    """ISO 8601 with the local offset, from epoch seconds (now if not given)."""
    if t is None:
        t = time.time()
    return _dt.datetime.fromtimestamp(t).astimezone().isoformat(timespec="seconds")


def parse_iso(s) -> float | None:
    """Epoch seconds of an ISO string written by `iso`, or None."""
    if not s or not isinstance(s, str):
        return None
    try:
        return _dt.datetime.fromisoformat(s).timestamp()
    except ValueError:
        return None


class StreamState:
    """What one `claude -p --output-format stream-json --verbose` stream told us so far."""

    def __init__(self):
        self.init = None              # the system/init event
        self.session_id = None
        self.msg_ids: list = []
        self._ids: set = set()
        self.first_usage = None       # usage of the first assistant event (cache verdict)
        self.first_at = None          # its arrival time (the base prefix read)
        self.last_usage = None        # usage of the last assistant event (base context)
        self.last_assistant_at = None
        self.result = None            # the result event
        self.refused = None           # text of a model-restricted warning
        self.events = 0

    def feed(self, ev: dict, t: float | None) -> None:
        if not isinstance(ev, dict):
            return
        self.events += 1
        ty = ev.get("type")
        if ty == "system":
            st = ev.get("subtype")
            if st == "init" and self.init is None:
                self.init = ev
                self.session_id = ev.get("session_id") or self.session_id
            elif (st == "informational" and ev.get("level") == "warning"
                  and "is restricted" in str(ev.get("content", ""))):
                self.refused = str(ev.get("content"))
        elif ty == "assistant":
            m = ev.get("message") or {}
            mid = m.get("id")
            if mid not in self._ids:
                self._ids.add(mid)
                self.msg_ids.append(mid)
            u = m.get("usage") or {}
            if self.first_usage is None:
                self.first_usage, self.first_at = u, t
            self.last_usage, self.last_assistant_at = u, t
        elif ty == "result":
            self.result = ev
            self.session_id = ev.get("session_id") or self.session_id

    @property
    def turns(self) -> int:
        return len(self.msg_ids)


def context_tokens(usage: dict | None) -> int:
    """Prompt-side tokens of one request: input + cache creation + cache read."""
    if not usage:
        return 0
    return (int(usage.get("input_tokens") or 0) + int(usage.get("cache_creation_input_tokens") or 0)
            + int(usage.get("cache_read_input_tokens") or 0))


_API_ERROR_TEXT = (   # checked in this order, after the status code
    ("prompt_too_long", re.compile(r"prompt is too long|prompt too long|prompt_too_long", re.I)),
    ("usage_limit", re.compile(r"usage limit|rate[ _-]?limit|hit your limit|api error:?\s*429", re.I)),
    ("overloaded", re.compile(r"overloaded|api error:?\s*529", re.I)),
    ("api_error", re.compile(r"api error", re.I)),
)


def api_error_reason(res: dict | None) -> tuple[str, str] | None:
    """(reason, detail) for a result that carries an API error, else None. `reason` is
    `usage_limit` (429), `overloaded` (529), `prompt_too_long` or `api_error`."""
    if not res:
        return None
    status = res.get("api_error_status")
    try:
        status = int(status) if status not in (None, "") else None
    except (TypeError, ValueError):
        status = None
    if not res.get("is_error") and status is None:
        return None
    parts = [res.get("result"), res.get("error"), res.get("errors")]
    text = " | ".join(p if isinstance(p, str) else json.dumps(p) for p in parts if p)
    detail = "api_error_status %s: %s" % (status, text[:200] if text else "(no text)")
    if status == 429:
        return "usage_limit", detail
    if status == 529:
        return "overloaded", detail
    for reason, rx in _API_ERROR_TEXT:
        if rx.search(text):
            return reason, detail
    if status is not None:
        return "api_error", detail
    return None


def classify(ss: StreamState, exit_code, stopped: bool) -> tuple[str, str | None]:
    """(status, reason) of a finished call; the first matching case wins."""
    res = ss.result
    match (bool(stopped), ss.refused, res, api_error_reason(res), exit_code):
        case (True, _, _, _, _):
            return "stopped", "stopped"
        case (_, str(), _, _, _):
            return "error", "model_refused"
        case (_, _, None, _, code):
            return "error", "exit_code" if code else "no_result"
        case (_, _, _, (reason, _), _):
            return "error", reason
        case (_, _, dict() as r, _, _) if r.get("is_error"):
            return "error", "is_error"
        case (_, _, _, _, code) if code:
            return "error", "exit_code"
        case _:
            return "ok", None


def usage_from_result(result: dict, before_mu: dict, cost_before) -> dict:
    """Tokens (with the 1h/5m split) from `result.usage`; `models_used` = modelUsage entries that
    grew against the session's totals before the call, largest first; cost delta against the
    session's cost before the call. Costs are rounded here, once, at 6 decimals."""
    result = result or {}
    usage = result.get("usage") or {}
    cc = usage.get("cache_creation") or {}
    tokens = {
        "input": int(usage.get("input_tokens") or 0),
        "output": int(usage.get("output_tokens") or 0),
        "cache_creation": int(usage.get("cache_creation_input_tokens") or 0),
        "cache_read": int(usage.get("cache_read_input_tokens") or 0),
        "cache_creation_1h": int(cc.get("ephemeral_1h_input_tokens") or 0),
        "cache_creation_5m": int(cc.get("ephemeral_5m_input_tokens") or 0),
    }
    after = result.get("modelUsage") or {}
    before_mu = before_mu or {}
    delta = {}
    for m, v in after.items():
        b = before_mu.get(m) or {}
        delta[m] = {f: int((v or {}).get(f) or 0) - int(b.get(f) or 0) for f in MU_FIELDS}
    grew = [m for m, d in delta.items() if any(x > 0 for x in d.values())]
    models_used = sorted(grew, key=lambda m: (-sum(delta[m].values()), m))
    sums = {f: sum(d[f] for d in delta.values()) for f in MU_FIELDS}
    mismatch = (sums["inputTokens"], sums["outputTokens"], sums["cacheCreationInputTokens"],
                sums["cacheReadInputTokens"]) != (tokens["input"], tokens["output"],
                                                   tokens["cache_creation"], tokens["cache_read"])
    total = result.get("total_cost_usd")
    if total is None:
        cost_delta = None
    else:
        cost_delta = round(float(total) - float(cost_before or 0.0), 6)
        total = round(float(total), 6)
    return {"tokens": tokens, "models_used": models_used, "mismatch": mismatch,
            "cost_delta": cost_delta, "total": total, "model_usage": after}


def cache_verdict(first_usage: dict | None, base_ctx) -> str | None:
    """The first request's cache read against the base's context tokens."""
    if first_usage is None:
        return None
    cr = int(first_usage.get("cache_read_input_tokens") or 0)
    if base_ctx and cr >= CACHE_HIT_RATIO * float(base_ctx):
        return "hit"
    return "partial" if cr > 0 else "miss"


def infer_ttl(tokens: dict | None, prev):
    """The TTL of a base's cache entries from one call's tokens: 1h written -> 3600, only 5m
    written -> 300, otherwise (read only, or nothing) `prev`."""
    if not tokens:
        return prev
    if int(tokens.get("cache_creation_1h") or 0) > 0:
        return TTL_1H
    if int(tokens.get("cache_creation_5m") or 0) > 0:
        return TTL_5M
    return prev


def wrote_only_5m(tokens: dict | None) -> bool:
    """The call wrote 5-minute cache entries and no 1-hour ones."""
    return (bool(tokens) and int(tokens.get("cache_creation_5m") or 0) > 0
            and int(tokens.get("cache_creation_1h") or 0) == 0)


def warmth(last_hit_ts: float | None, ttl_s, now: float) -> dict:
    """warm until `margin` before last_hit + ttl, expiring until last_hit + ttl, then cold;
    margin = max(120, ttl/10). No hit at all is cold."""
    if last_hit_ts is None:
        return {"base_state": "cold", "warm_until": None, "ttl_left_s": None}
    ttl = float(ttl_s or TTL_1H)
    margin = max(120.0, ttl / 10.0)
    wu = float(last_hit_ts) + ttl
    state = "warm" if now < wu - margin else ("expiring" if now < wu else "cold")
    return {"base_state": state, "warm_until": iso(wu), "ttl_left_s": int(max(0.0, wu - now))}