measurement.py

"""Measurement / "spend avoided" rollup.

Every routed task is logged as one JSON line in an append-only usage log, and ``tanglebrain
--stats`` rolls those records up into a "spend avoided" figure: what the routed work *would* have
cost on a paid frontier API, had it not gone to the free local tier or a subscription CLI. This
makes the cloud-equivalent cost avoided by routing visible.

Design notes:

- **Tokens are estimated, not measured.** Authenticated CLIs expose no usable token counts, and a
  local reasoning model's real ``usage`` is inflated by dropped reasoning tokens. So a
  single ``chars/4`` heuristic over the visible prompt + response is applied *uniformly* across all
  tiers — one consistent, if approximate, methodology (see :func:`estimate_tokens`).
- **Pricing is config-driven** (``config/pricing.yaml``): a reference frontier price the operator
  tunes (the knob GUI edits it).
- **All I/O is fault-tolerant.** A logging failure must never break the user's actual answer, and a
  corrupt log line must never break the rollup. Reads return sensible defaults; the writer swallows
  every exception. This mirrors the router's state-file idiom (:mod:`tanglebrain.router`).
"""
from __future__ import annotations

import json
import os
import shutil
import threading
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import yaml

from tanglebrain.router import DEFAULT_STATE_SUBDIR, STATE_DIR_ENV

LOG_FILENAME = "usage.jsonl"

#: Env var carrying the top-level task id from the orchestrator down to a delegated sub-call. The
#: CLI mints a task id per routed task and the orchestrator-CLI adapter injects it into the
#: orchestrator subprocess env (only when the delegate tool is injected); the orchestrator forwards
#: its env to the MCP delegate child it spawns, where :func:`tanglebrain.delegate.run_delegate` reads
#: it back and stamps each delegate record's ``parent_task_id``. This is what links a delegated
#: sub-call to the specific top-level task that spawned it, across the process boundary.
PARENT_TASK_ID_ENV = "TANGLEBRAIN_TASK_ID"

# Serializes appends to the usage log so concurrent writers in one process (delegate_many fans
# sub-tasks out across threads) can't interleave bytes mid-line. Per-process only; cross-process
# appends rely on the OS's O_APPEND atomicity for short lines, as before.
_LOG_LOCK = threading.Lock()

# Chars per token for the uniform estimation heuristic. ~4 chars/token is the standard rough
# approximation for English-ish text across modern BPE tokenizers; good enough for an *estimate*.
_CHARS_PER_TOKEN = 4


@dataclass(frozen=True)
class Pricing:
    """Cloud-equivalent reference pricing for the rollup (loaded from ``config/pricing.yaml``).

    Attributes:
        reference_model: Human-readable label for the frontier model these rates represent.
        input_per_mtok: US dollars per 1,000,000 input (prompt) tokens.
        output_per_mtok: US dollars per 1,000,000 output (completion) tokens.
        is_placeholder: ``True`` while the rates are rough/illustrative; the rollup renders a
            PLACEHOLDER caveat so no figure is mistaken for a precise cost.
    """

    reference_model: str
    input_per_mtok: float
    output_per_mtok: float
    is_placeholder: bool


# Fallback used when ``config/pricing.yaml`` is missing or unreadable — always flagged placeholder.
PLACEHOLDER_PRICING = Pricing(
    reference_model="unconfigured (PLACEHOLDER — pricing.yaml unreadable)",
    input_per_mtok=3.00,
    output_per_mtok=15.00,
    is_placeholder=True,
)


def default_log_path() -> Path:
    """Return the usage-log file path.

    Honors ``TANGLEBRAIN_STATE_DIR`` (``~`` expanded); otherwise ``~/.cache/tanglebrain/``. The log
    lives alongside the router's state file (same dir, same env override).

    Returns:
        The absolute path to the append-only usage JSONL file.
    """
    base = os.environ.get(STATE_DIR_ENV)
    root = Path(base).expanduser() if base else Path.home() / DEFAULT_STATE_SUBDIR
    return root / LOG_FILENAME


def default_pricing_path() -> Path:
    """Return the path to the pricing YAML shipped with the package.

    Returns:
        The absolute path to ``tanglebrain/config/pricing.yaml``.
    """
    return Path(__file__).resolve().parent / "config" / "pricing.yaml"


def load_pricing(path: str | os.PathLike[str] | None = None) -> Pricing:
    """Load cloud-equivalent reference pricing, tolerating a missing/corrupt file.

    Args:
        path: Path to a pricing YAML. Defaults to the packaged ``config/pricing.yaml``.

    Returns:
        The parsed :class:`Pricing`, or :data:`PLACEHOLDER_PRICING` if the file is absent,
        unreadable, or malformed — bad config must never crash the rollup.
    """
    pricing_path = Path(path) if path is not None else default_pricing_path()
    try:
        raw = yaml.safe_load(pricing_path.read_text())
        return Pricing(
            reference_model=str(raw.get("reference_model", "unknown")),
            input_per_mtok=float(raw["input_per_mtok"]),
            output_per_mtok=float(raw["output_per_mtok"]),
            is_placeholder=bool(raw.get("placeholder", False)),
        )
    except (OSError, yaml.YAMLError, ValueError, TypeError, KeyError, AttributeError):
        return PLACEHOLDER_PRICING


# Fallback header, used ONLY when the target file is absent (e.g. a fresh write to a new path).
# A normal save preserves the existing file's own leading comment block verbatim (see
# :func:`_leading_comment_block`), so the curated methodology note is never replaced or drifted.
PRICING_HEADER = """\
# Cloud-equivalent reference pricing for the "spend avoided" rollup.
#
# Methodology: for each routed task, estimate what it WOULD have cost on a paid frontier API, valued
# at a reference model's per-million-token price. The rollup multiplies estimated tokens (a chars/4
# heuristic over the visible prompt + response) by these rates. Values are US dollars per 1,000,000
# tokens.
#
# `placeholder: true` makes `tanglebrain --stats` flag every figure as PLACEHOLDER (use it if you
# fork these reference rates before re-checking them). Edited via `tanglebrain-gui` or by hand.
"""


def _leading_comment_block(text: str) -> str:
    """Return the file's leading run of comment/blank lines (its header), or ``""`` if none.

    Used to preserve a pricing file's curated header verbatim across a save, so no documentation
    is lost or replaced. Stops at the first non-comment, non-blank line (the first YAML key).
    """
    out: list[str] = []
    for line in text.splitlines():
        if line.startswith("#") or not line.strip():
            out.append(line)
        else:
            break
    while out and not out[-1].strip():  # drop trailing blank lines before the keys
        out.pop()
    return "\n".join(out) + "\n" if out else ""


def validate_pricing(data: dict) -> Pricing:
    """Strictly validate raw pricing fields and build a :class:`Pricing`.

    Unlike :func:`load_pricing` (lenient — bad reads fall back to a placeholder), this rejects
    invalid input so the panel never persists garbage.

    Args:
        data: ``{reference_model, input_per_mtok, output_per_mtok, placeholder}``.

    Returns:
        A validated :class:`Pricing`.

    Raises:
        ValueError: If a field is missing, the wrong type, a non-finite/negative rate, or an
            empty ``reference_model``.
    """
    if not isinstance(data, dict):
        raise ValueError("pricing must be an object")

    model = data.get("reference_model")
    if not isinstance(model, str) or not model.strip():
        raise ValueError("reference_model must be a non-empty string")

    rates = {}
    for key in ("input_per_mtok", "output_per_mtok"):
        value = data.get(key)
        if isinstance(value, bool) or not isinstance(value, (int, float)):
            raise ValueError(f"{key} must be a number")
        value = float(value)
        if value != value or value in (float("inf"), float("-inf")):  # NaN / inf guard
            raise ValueError(f"{key} must be finite")
        if value < 0:
            raise ValueError(f"{key} must be >= 0")
        rates[key] = value

    placeholder = data.get("placeholder", False)
    if not isinstance(placeholder, bool):
        raise ValueError("placeholder must be a boolean")

    return Pricing(
        reference_model=model.strip(),
        input_per_mtok=rates["input_per_mtok"],
        output_per_mtok=rates["output_per_mtok"],
        is_placeholder=placeholder,
    )


def _backup_dir() -> Path:
    """Return the directory for config backups (under the state dir, never the repo config dir)."""
    base = os.environ.get(STATE_DIR_ENV)
    root = Path(base).expanduser() if base else Path.home() / DEFAULT_STATE_SUBDIR
    return root / "backups"


def _atomic_write(path: Path, text: str) -> None:
    """Write ``text`` to ``path`` atomically (temp file in the same dir, then ``os.replace``)."""
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + ".tmp")
    tmp.write_text(text, encoding="utf-8")
    os.replace(tmp, path)


def _render_pricing(pricing: Pricing, header: str) -> str:
    """Render a :class:`Pricing` to YAML text beneath ``header``.

    ``reference_model`` is emitted via ``json.dumps`` — a JSON string is valid YAML and safely
    quotes/escapes any colons, quotes, unicode, or backslashes. Float rates use Python's ``repr``
    (valid YAML), which round-trips exactly through :func:`load_pricing`.
    """
    return (
        header
        + f"placeholder: {str(pricing.is_placeholder).lower()}\n"
        + f"reference_model: {json.dumps(pricing.reference_model)}\n"
        + f"input_per_mtok: {pricing.input_per_mtok}\n"
        + f"output_per_mtok: {pricing.output_per_mtok}\n"
    )


def save_pricing(pricing: Pricing, path: str | os.PathLike[str] | None = None) -> None:
    """Persist pricing to the config YAML — header-preserving, with a backup, written atomically.

    Preserves the target's existing leading comment block verbatim (falling back to
    :data:`PRICING_HEADER` only when the file is absent), backs up any existing file to
    ``<state_dir>/backups/pricing-<ts>.yaml``, then atomically replaces the target.

    Args:
        pricing: The validated pricing to write (see :func:`validate_pricing`).
        path: Target YAML path. Defaults to the packaged ``config/pricing.yaml``.
    """
    target = Path(path) if path is not None else default_pricing_path()
    header = PRICING_HEADER
    if target.exists():
        existing = target.read_text(encoding="utf-8")
        block = _leading_comment_block(existing)
        if block.strip():
            header = block  # keep the curated header verbatim — no drift, no doc loss
        backup_dir = _backup_dir()
        backup_dir.mkdir(parents=True, exist_ok=True)
        stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S_%fZ")  # sub-second: no same-second collision
        shutil.copy2(target, backup_dir / f"pricing-{stamp}.yaml")
    _atomic_write(target, _render_pricing(pricing, header))


def estimate_tokens(text: str) -> int:
    """Estimate the token count of ``text`` via the uniform ``chars/4`` heuristic.

    This is an approximation, applied identically to every tier (CLI subs expose no real counts).
    Empty/falsy text counts as 0; any non-empty text is at least 1 token.

    Args:
        text: The prompt or response text.

    Returns:
        The estimated token count (``>= 0``).
    """
    if not text:
        return 0
    return max(1, len(text) // _CHARS_PER_TOKEN)


def cloud_equiv_usd(in_tokens: int, out_tokens: int, pricing: Pricing) -> float:
    """Compute the cloud-equivalent cost of a task at the reference frontier price.

    Args:
        in_tokens: Estimated input (prompt) tokens.
        out_tokens: Estimated output (completion) tokens.
        pricing: The reference pricing to apply.

    Returns:
        The estimated US-dollar cost on the reference frontier API.
    """
    return (
        in_tokens / 1_000_000 * pricing.input_per_mtok
        + out_tokens / 1_000_000 * pricing.output_per_mtok
    )


def record_task(
    *,
    path: str,
    entry: object,
    prompt: str,
    response: str,
    kind: str = "task",
    task_id: str | None = None,
    parent_task_id: str | None = None,
    origin: str | None = None,
    log_path: str | os.PathLike[str] | None = None,
    pricing: Pricing | None = None,
) -> None:
    """Append one usage record for a routed task or a delegated sub-call. Never raises.

    A logging failure is dropped — measurement is a side-effect that must never affect the returned
    answer.

    Args:
        path: Which execution path served the work — ``router`` | ``local`` | ``model`` |
            ``delegate``.
        entry: The served :class:`~tanglebrain.roster.RosterEntry` (read for ``tier``/``id``); may
            be ``None`` (e.g. the router didn't surface one), in which case both are ``"unknown"``.
        prompt: The prompt (for input-token estimation).
        response: The returned response text (for output-token estimation).
        kind: ``"task"`` for a top-level routed task (the default; what the spend-avoided headline
            counts) or ``"delegate"`` for a delegated sub-call. Delegate records are rolled up
            **separately** so a sub-call's saving is never double-counted against its parent task.
        task_id: For a top-level task, the id minted for this routed task (so its delegated sub-calls
            can be linked back to it). Omitted from the record when ``None``.
        parent_task_id: For a delegated sub-call, the id of the top-level task that spawned it (read
            from :data:`PARENT_TASK_ID_ENV`). Omitted from the record when ``None`` — e.g. a delegate
            invoked outside a propagated task, which rolls up as ``unlinked``. For a top-level task,
            an external caller's own task/session identity (#74: the serve endpoint's
            ``X-TangleBrain-Parent-Task`` header) — pure attribution metadata; the delegate tree's
            ``by_parent`` rollup reads it only off ``delegate`` records.
        origin: Which surface the work entered through — ``"cli"`` | ``"gui"`` | ``"serve"``
            (#74). Omitted from the record when ``None``; records without it roll up as
            ``untagged`` (pre-#74 history is never guessed at).
        log_path: Override the usage-log path (tests inject a temp path). Defaults to
            :func:`default_log_path`.
        pricing: Override the pricing. Defaults to :func:`load_pricing`.
    """
    try:
        if pricing is None:
            pricing = load_pricing()
        tier = getattr(entry, "tier", None) or "unknown"
        model = getattr(entry, "id", None) or "unknown"
        in_tok = estimate_tokens(prompt)
        out_tok = estimate_tokens(response)
        equiv = cloud_equiv_usd(in_tok, out_tok, pricing)
        # A paid `api` task incurs real spend, so it avoids nothing (avoided = 0); every other tier
        # routes work off a paid frontier API, so it avoids the full cloud-equivalent.
        avoided = 0.0 if tier == "api" else equiv
        record = {
            "ts": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "kind": str(kind),
            "path": str(path),
            "tier": str(tier),
            "model": str(model),
            "in_tokens_est": in_tok,
            "out_tokens_est": out_tok,
            "cloud_equiv_usd": round(equiv, 6),
            "spend_avoided_usd": round(avoided, 6),
            "pricing_ref": pricing.reference_model,
        }
        # Optional linkage fields — only written when present, so existing records/readers that
        # never set them are unaffected (a missing field reads as "no linkage").
        if task_id is not None:
            record["task_id"] = str(task_id)
        if parent_task_id is not None:
            record["parent_task_id"] = str(parent_task_id)
        if origin is not None:
            record["origin"] = str(origin)
        target = Path(log_path) if log_path is not None else default_log_path()
        target.parent.mkdir(parents=True, exist_ok=True)
        with _LOG_LOCK:
            with target.open("a", encoding="utf-8") as fh:
                fh.write(json.dumps(record) + "\n")
    except Exception:
        # Measurement is a side-effect: a failure here must never affect the returned answer.
        return


def _as_int(value: object) -> int:
    """Coerce a stored numeric field to int, defaulting to 0 on any bad value."""
    try:
        return int(value)  # type: ignore[arg-type]
    except (ValueError, TypeError):
        return 0


def _as_float(value: object) -> float:
    """Coerce a stored numeric field to float, defaulting to 0.0 on any bad value."""
    try:
        return float(value)  # type: ignore[arg-type]
    except (ValueError, TypeError):
        return 0.0


def read_records(log_path: str | os.PathLike[str] | None = None) -> list[dict]:
    """Read all usage records from the log, skipping malformed lines.

    Args:
        log_path: Override the usage-log path. Defaults to :func:`default_log_path`.

    Returns:
        The parsed records in file (chronological) order. An absent log yields ``[]``.
    """
    target = Path(log_path) if log_path is not None else default_log_path()
    records: list[dict] = []
    try:
        text = target.read_text(encoding="utf-8")
    except OSError:
        return records
    for line in text.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(obj, dict):
            records.append(obj)
    return records


def rollup(records: list[dict]) -> dict:
    """Aggregate usage records into a summary.

    Args:
        records: The records from :func:`read_records`.

    Returns:
        A dict with: ``tasks`` (int), ``by_tier`` (tier → count), ``by_origin`` (origin → count,
        where a record without an ``origin`` field counts as ``untagged`` — pre-#74 history is
        never guessed at), ``in_tokens_est`` / ``out_tokens_est`` (summed estimates), and
        ``cloud_equiv_usd`` / ``spend_avoided_usd``
        (summed dollars) — all over **top-level tasks only** — plus ``delegates``, a separate
        sub-rollup of delegated sub-calls ``{count, by_backend: {model: {count, in_tokens_est,
        out_tokens_est}}, by_parent: {parent_task_id: {count, by_backend: {model: count}}},
        in_tokens_est, out_tokens_est, cloud_equiv_usd}``. ``by_parent`` groups each delegate under
        the top-level task that spawned it (via ``parent_task_id``); delegates with no
        ``parent_task_id`` are grouped under the sentinel ``"unlinked"``. Delegate records are kept
        out of the headline so a sub-call's saving is never double-counted against its parent task;
        their cloud-equiv is informational. A record without a ``kind`` field counts as a task.
    """
    summary: dict = {
        "tasks": 0,
        "by_tier": {},
        "by_origin": {},
        "in_tokens_est": 0,
        "out_tokens_est": 0,
        "cloud_equiv_usd": 0.0,
        "spend_avoided_usd": 0.0,
    }
    delegates: dict = {
        "count": 0,
        "by_backend": {},
        "by_parent": {},
        "in_tokens_est": 0,
        "out_tokens_est": 0,
        "cloud_equiv_usd": 0.0,
    }
    for r in records:
        in_tok = _as_int(r.get("in_tokens_est"))
        out_tok = _as_int(r.get("out_tokens_est"))
        if str(r.get("kind", "task")) == "delegate":
            delegates["count"] += 1
            model = str(r.get("model", "unknown"))
            backend = delegates["by_backend"].setdefault(
                model, {"count": 0, "in_tokens_est": 0, "out_tokens_est": 0}
            )
            backend["count"] += 1
            backend["in_tokens_est"] += in_tok
            backend["out_tokens_est"] += out_tok
            # Per-parent tree: link this sub-call to the top-level task that spawned it. A delegate
            # with no parent_task_id (run outside a propagated task) groups under "unlinked".
            parent_id = r.get("parent_task_id")
            parent_key = str(parent_id) if parent_id not in (None, "") else "unlinked"
            parent = delegates["by_parent"].setdefault(parent_key, {"count": 0, "by_backend": {}})
            parent["count"] += 1
            parent["by_backend"][model] = parent["by_backend"].get(model, 0) + 1
            delegates["in_tokens_est"] += in_tok
            delegates["out_tokens_est"] += out_tok
            delegates["cloud_equiv_usd"] += _as_float(r.get("cloud_equiv_usd"))
            continue
        summary["tasks"] += 1
        tier = str(r.get("tier", "unknown"))
        summary["by_tier"][tier] = summary["by_tier"].get(tier, 0) + 1
        origin = str(r.get("origin") or "untagged")
        summary["by_origin"][origin] = summary["by_origin"].get(origin, 0) + 1
        summary["in_tokens_est"] += in_tok
        summary["out_tokens_est"] += out_tok
        summary["cloud_equiv_usd"] += _as_float(r.get("cloud_equiv_usd"))
        summary["spend_avoided_usd"] += _as_float(r.get("spend_avoided_usd"))
    summary["cloud_equiv_usd"] = round(summary["cloud_equiv_usd"], 4)
    summary["spend_avoided_usd"] = round(summary["spend_avoided_usd"], 4)
    delegates["cloud_equiv_usd"] = round(delegates["cloud_equiv_usd"], 4)
    summary["delegates"] = delegates
    return summary


def format_rollup(summary: dict, pricing: Pricing) -> str:
    """Render a rollup summary as a human-readable block for the CLI.

    Args:
        summary: The aggregate from :func:`rollup`.
        pricing: The currently-configured pricing (for the reference-model label + placeholder
            caveat). Per-record costs were computed when each task ran; this only labels the figure.

    Returns:
        A multi-line string suitable for printing.
    """
    lines = [
        "TangleBrain — spend avoided (cloud-equivalent)",
        f"  Tasks routed:   {summary.get('tasks', 0)}",
    ]
    by_tier = summary.get("by_tier") or {}
    if by_tier:
        tiers = ", ".join(f"{k} {v}" for k, v in sorted(by_tier.items()))
        lines.append(f"  By tier:        {tiers}")
    by_origin = summary.get("by_origin") or {}
    # Show the origin split only once it says something — all-untagged history adds no signal.
    if any(k != "untagged" for k in by_origin):
        origins = ", ".join(f"{k} {v}" for k, v in sorted(by_origin.items()))
        lines.append(f"  By origin:      {origins}")
    lines.append(
        f"  Est. tokens:    in {summary.get('in_tokens_est', 0):,} / "
        f"out {summary.get('out_tokens_est', 0):,}"
    )
    lines.append(f"  Spend avoided:  ${summary.get('spend_avoided_usd', 0.0):,.2f}")
    lines.append(f"  Pricing ref:    {pricing.reference_model}")
    if pricing.is_placeholder:
        lines.append(
            "  ⚠ pricing: PLACEHOLDER — figures are illustrative; set real rates in "
            "config/pricing.yaml and flip placeholder to false."
        )

    delegates = summary.get("delegates") or {}
    if delegates.get("count"):
        lines.append("")
        lines.append("  Delegated sub-tasks (offloaded by orchestrators)")
        lines.append(f"    Count:        {delegates.get('count', 0)}")
        by_backend = delegates.get("by_backend") or {}
        if by_backend:
            backends = ", ".join(
                f"{model} {info.get('count', 0)}" for model, info in sorted(by_backend.items())
            )
            lines.append(f"    By backend:   {backends}")
        by_parent = delegates.get("by_parent") or {}
        if by_parent:
            linked = [k for k in by_parent if k != "unlinked"]
            unlinked = (by_parent.get("unlinked") or {}).get("count", 0)
            if linked:
                tree = f"{len(linked)} parent task(s)"
                if unlinked:
                    tree += f", {unlinked} unlinked"
            else:
                tree = f"{unlinked} unlinked"  # all sub-calls ran outside a propagated task
            lines.append(f"    Linked to:    {tree}")
        lines.append(
            f"    Est. tokens:  in {delegates.get('in_tokens_est', 0):,} / "
            f"out {delegates.get('out_tokens_est', 0):,}"
        )
        lines.append(
            f"    Cloud-equiv:  ${delegates.get('cloud_equiv_usd', 0.0):,.2f} "
            "(informational — already credited within parent tasks)"
        )
    return "\n".join(lines)