#!/usr/bin/env python3
"""
dsh-cache-audit — what your DeepSeek prompt cache is actually worth.

DeepSeek bills input tokens on two meters: cache hit and cache miss. The gap is
large enough that your real bill is decided by a number neither the pricing page
nor your invoice shows you — the share of your input tokens that hit the cache.

DeepSeek Harness records that number for every single request. This script reads
those records off your own disk and turns them into a bill.

    python3 dsh-cache-audit.py

Nothing is uploaded. Nothing is sent anywhere. The script reads session logs,
adds up token counters, and prints. It never opens ~/.dsh/.credentials.yaml and
never touches the network — read it top to bottom before running it, it is one
file with no dependencies.

Requires: Python 3.9+, and a way to read zstd (any one of: Python 3.14's
built-in `compression.zstd`, the `zstandard` package, or the `zstd` CLI).

Prices verified against api-docs.deepseek.com/quick_start/pricing on 2026-08-14.
More context, and the numbers from a 157M-token run, at:
https://deepseekprice.com/deepseek-cache-pricing
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path

# --------------------------------------------------------------------------
# Rate cards, US dollars per million tokens.
#
# The old card was flat — one price around the clock. The new one splits by UTC
# hour, so the same request costs two different amounts depending on when it is
# sent. Both are kept here because the interesting question is not what you will
# pay, it is what the change did to a bill shaped like yours.
# --------------------------------------------------------------------------

PRICE_CHANGE_UTC = "2026-08-16T16:00:00Z"

# (cache hit, cache miss, output)
RATES = {
    "deepseek-v4-pro": {
        "old": (0.003625, 0.435, 0.87),
        "off_peak": (0.022, 0.66, 1.98),
        "peak": (0.044, 1.32, 3.96),
    },
    "deepseek-v4-flash": {
        "old": (0.0028, 0.14, 0.28),
        "off_peak": (0.007, 0.22, 0.66),
        "peak": (0.014, 0.44, 1.32),
    },
}

# Peak windows on the UTC clock, as (start_hour, end_hour). Everything else is
# off-peak. Note where these land: 01:00-04:00 and 06:00-10:00 UTC is the
# middle of the working day in Beijing and the middle of the night in New York.
PEAK_WINDOWS_UTC = ((1, 4), (6, 10))

MILLION = 1_000_000


def is_peak(when: datetime) -> bool:
    hour = when.astimezone(timezone.utc).hour
    return any(start <= hour < end for start, end in PEAK_WINDOWS_UTC)


# --------------------------------------------------------------------------
# Reading the logs
# --------------------------------------------------------------------------


def decompress(path: Path) -> str:
    """Return the decompressed contents of a .zstd file.

    Three ways to do this and no way to rely on any one of them: the stdlib only
    grew zstd in 3.14, the `zstandard` package is not installed by default, and
    the CLI may not be on PATH. Try each, complain usefully if all three fail.
    """
    raw = path.read_bytes()

    try:
        from compression import zstd  # type: ignore[import-not-found]  # Python 3.14+

        return zstd.decompress(raw).decode("utf-8", "replace")
    except ImportError:
        pass

    try:
        import zstandard  # type: ignore[import-not-found]

        return zstandard.ZstdDecompressor().stream_reader(raw).read().decode("utf-8", "replace")
    except ImportError:
        pass

    if shutil.which("zstd"):
        done = subprocess.run(["zstd", "-dc", str(path)], capture_output=True)
        if done.returncode == 0:
            return done.stdout.decode("utf-8", "replace")

    sys.exit(
        "Cannot read zstd files. Install one of:\n"
        "  pip install zstandard\n"
        "  brew install zstd      (macOS)\n"
        "  apt install zstd       (Debian/Ubuntu)"
    )


class Call:
    """One request to the model, as recorded by dsh.

    dsh reports usage twice — streaming deltas on `assistant/chunk`, then a final
    figure on `assistant/message`. Only the latter is counted here; summing the
    chunks double-counts every request.

    On the field names: `inputTokens` counts input that did NOT hit the cache,
    and `cacheReadTokens` counts input that did. They are disjoint, so total
    prompt size is the sum of the two. (Verified empirically: across 650 real
    requests `cacheReadTokens` exceeded `inputTokens` 631 times and equalled it
    never, which rules out the other reading, where one contains the other.)
    `reasoningTokens` is a subset of `outputTokens` — thinking tokens you are
    billed for at the output rate but never see.
    """

    __slots__ = ("session", "model", "index", "at", "miss", "hit", "output", "reasoning")

    def __init__(self, session, model, index, at, usage):
        self.session = session
        self.model = model or "unknown"
        self.index = index
        self.at = at
        self.miss = usage.get("inputTokens", 0) or 0
        self.hit = usage.get("cacheReadTokens", 0) or 0
        self.output = usage.get("outputTokens", 0) or 0
        self.reasoning = usage.get("reasoningTokens", 0) or 0

    @property
    def prompt(self) -> int:
        return self.hit + self.miss


def read_session(path: Path) -> list[Call]:
    """Parse one session.jsonl.zstd into its model calls, in order."""
    calls: list[Call] = []
    model = None
    for line in decompress(path).splitlines():
        if not line.strip():
            continue
        try:
            event = json.loads(line)
        except json.JSONDecodeError:
            continue  # a session being written to can end mid-line

        kind, data = event.get("type"), event.get("data") or {}

        # The model is announced once per request, not carried on the usage
        # record — so track the most recent announcement and attribute to it.
        if kind == "request/header":
            model = ((data.get("header") or {}).get("config") or {}).get("model")

        elif kind == "assistant/message" and isinstance(data.get("usage"), dict):
            stamp = event.get("time")
            at = (
                datetime.fromtimestamp(stamp / 1000, timezone.utc)
                if isinstance(stamp, (int, float))
                else None
            )
            calls.append(Call(path.parent.name, model, len(calls) + 1, at, data["usage"]))
    return calls


# --------------------------------------------------------------------------
# Costing
# --------------------------------------------------------------------------


def cost(hit: int, miss: int, output: int, rates: tuple[float, float, float]) -> float:
    hit_rate, miss_rate, out_rate = rates
    return (hit * hit_rate + miss * miss_rate + output * out_rate) / MILLION


def summarise(calls: list[Call]) -> dict:
    hit = sum(c.hit for c in calls)
    miss = sum(c.miss for c in calls)
    output = sum(c.output for c in calls)
    reasoning = sum(c.reasoning for c in calls)
    prompt = hit + miss

    summary = {
        "calls": len(calls),
        "prompt_tokens": prompt,
        "cache_hit_tokens": hit,
        "cache_miss_tokens": miss,
        "cache_hit_rate": hit / prompt if prompt else 0.0,
        "output_tokens": output,
        "reasoning_tokens": reasoning,
        "reasoning_share_of_output": reasoning / output if output else 0.0,
    }

    model = calls[0].model if calls else "unknown"
    cards = RATES.get(model)
    if cards:
        bills = {name: cost(hit, miss, output, card) for name, card in cards.items()}
        # The counterfactual that makes the cache legible: identical work, no
        # cache. Every hit token is repriced as a miss.
        no_cache = {name: cost(0, prompt, output, card) for name, card in cards.items()}
        summary["bill"] = bills
        summary["bill_without_cache"] = no_cache
        summary["cache_saved_old"] = no_cache["old"] - bills["old"]
        summary["cache_saved_off_peak"] = no_cache["off_peak"] - bills["off_peak"]
        summary["increase_actual"] = bills["off_peak"] / bills["old"] if bills["old"] else 0
        summary["increase_without_cache"] = (
            no_cache["off_peak"] / no_cache["old"] if no_cache["old"] else 0
        )
    return summary


def peak_split(calls: list[Call]) -> dict:
    """How much of the work landed inside the expensive windows."""
    timed = [c for c in calls if c.at]
    if not timed:
        return {}
    in_peak = [c for c in timed if is_peak(c.at)]
    peak_tokens = sum(c.prompt + c.output for c in in_peak)
    all_tokens = sum(c.prompt + c.output for c in timed)
    return {
        "calls_total": len(timed),
        "calls_in_peak": len(in_peak),
        "token_share_in_peak": peak_tokens / all_tokens if all_tokens else 0.0,
        "by_utc_hour": {
            hour: sum(1 for c in timed if c.at.astimezone(timezone.utc).hour == hour)
            for hour in range(24)
        },
    }


# Buckets by absolute call number, not by percentage through the session. The
# warm-up is over within a handful of calls, so percentage buckets bury it: on a
# 300-call session the first bucket is 50 calls and the one cold request gets
# averaged into invisibility.
WARMUP_BUCKETS = ((1, 1), (2, 2), (3, 3), (4, 5), (6, 10), (11, None))


def warmup_curve(sessions: dict[str, list[Call]]) -> list[dict]:
    """Cache hit rate as a function of how deep into a session you are.

    A cold session has nothing to hit — the first request pays full price for
    every token. The question worth answering is how many requests it takes for
    that to stop being true.
    """
    curve = []
    for low, high in WARMUP_BUCKETS:
        bucket = [
            c
            for calls in sessions.values()
            for c in calls
            if c.index >= low and (high is None or c.index <= high)
        ]
        prompt = sum(c.prompt for c in bucket)
        curve.append(
            {
                "position": f"#{low}" if low == high else f"#{low}-{high}" if high else f"#{low}+",
                "calls": len(bucket),
                "cache_hit_rate": (sum(c.hit for c in bucket) / prompt) if prompt else 0.0,
            }
        )
    return curve


# --------------------------------------------------------------------------
# Output
# --------------------------------------------------------------------------


def money(value: float) -> str:
    return f"${value:,.4f}" if value < 10 else f"${value:,.2f}"


def report(sessions: dict[str, list[Call]], by_model: dict[str, list[Call]]) -> None:
    everything = [c for calls in sessions.values() for c in calls]
    print()
    print("  DeepSeek cache audit")
    print(f"  {len(everything):,} model calls across {len(sessions)} sessions")
    print("  " + "─" * 68)

    for model, calls in sorted(by_model.items(), key=lambda kv: -len(kv[1])):
        s = summarise(calls)
        print(f"\n  {model} — {s['calls']:,} calls")
        print(
            f"    prompt   {s['prompt_tokens']:>15,}  "
            f"= {s['cache_hit_tokens']:,} cached + {s['cache_miss_tokens']:,} uncached"
        )
        print(f"    cache hit rate           {s['cache_hit_rate']:>8.2%}")
        print(
            f"    output   {s['output_tokens']:>15,}  "
            f"({s['reasoning_share_of_output']:.1%} of it invisible thinking, billed in full)"
        )

        if "bill" not in s:
            print("    (no rate card on file for this model — tokens only)")
            continue

        b, nc = s["bill"], s["bill_without_cache"]
        print()
        print(f"    {'':<22}{'as billed':>14}{'if 0% cached':>16}{'cache saved':>14}")
        for label, key in (("before 16 Aug", "old"), ("off-peak now", "off_peak"), ("peak now", "peak")):
            saved = nc[key] - b[key]
            print(f"    {label:<22}{money(b[key]):>14}{money(nc[key]):>16}{money(saved):>14}")

        print()
        print(f"    The change cost a bill shaped like yours {s['increase_actual']:.2f}x more.")
        print(f"    The same work with no cache would have gone up {s['increase_without_cache']:.2f}x.")
        if s["increase_actual"] > s["increase_without_cache"]:
            hit_multiple = RATES[model]["off_peak"][0] / RATES[model]["old"][0]
            print(
                f"    You are hit harder because you cache well: the cache-hit rate rose"
                f" {hit_multiple:.2f}x,\n    the steepest rise of the three meters."
            )

    peaks = peak_split(everything)
    if peaks:
        print("\n  " + "─" * 68)
        print(
            f"\n  Timing: {peaks['calls_in_peak']:,} of {peaks['calls_total']:,} calls"
            f" ({peaks['token_share_in_peak']:.1%} of tokens) landed in peak hours."
        )
        offset = datetime.now().astimezone().utcoffset()
        shift = int(offset.total_seconds() // 3600) if offset else 0
        print(f"  Peak is 01:00-04:00 and 06:00-10:00 UTC. Your clock runs UTC{shift:+d}.\n")
        busiest = max(peaks["by_utc_hour"].values())
        for hour, count in sorted(peaks["by_utc_hour"].items()):
            if not count:
                continue
            bar = "█" * max(1, count * 36 // busiest)
            tag = "  PEAK" if is_peak(datetime(2026, 1, 1, hour, tzinfo=timezone.utc)) else ""
            print(
                f"    {hour:02d}:00 UTC  ({(hour + shift) % 24:02d}:00 local)"
                f"  {count:>5,} calls  {bar}{tag}"
            )

    curve = warmup_curve(sessions)
    if curve:
        print("\n  " + "─" * 68)
        print("\n  Warm-up: cache hit rate by call number within a session")
        for point in curve:
            bar = "█" * int(point["cache_hit_rate"] * 36)
            print(
                f"    call {point['position']:>6}  {point['cache_hit_rate']:>7.1%}"
                f"  {bar}  ({point['calls']:,} calls)"
            )
    print()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compute your real DeepSeek cache hit rate and bill from dsh session logs.",
    )
    parser.add_argument(
        "--sessions",
        type=Path,
        default=Path.home() / ".dsh" / "sessions",
        help="dsh sessions directory (default: ~/.dsh/sessions)",
    )
    parser.add_argument("--json", action="store_true", help="emit machine-readable JSON instead")
    args = parser.parse_args()

    if not args.sessions.is_dir():
        sys.exit(f"No session directory at {args.sessions}. Point --sessions at yours.")

    logs = sorted(args.sessions.rglob("session.jsonl.zstd"))
    if not logs:
        sys.exit(f"No session logs under {args.sessions}.")

    sessions: dict[str, list[Call]] = {}
    by_model: dict[str, list[Call]] = defaultdict(list)
    for log in logs:
        calls = read_session(log)
        if not calls:
            continue  # a session that never reached the model
        sessions[log.parent.name] = calls
        for call in calls:
            by_model[call.model].append(call)

    if not sessions:
        sys.exit("Found session logs, but none of them contain model calls yet.")

    if args.json:
        print(
            json.dumps(
                {
                    "sessions": len(sessions),
                    "price_change_utc": PRICE_CHANGE_UTC,
                    "by_model": {m: summarise(c) for m, c in by_model.items()},
                    "peak_split": peak_split([c for cs in sessions.values() for c in cs]),
                    "warmup_curve": warmup_curve(sessions),
                },
                indent=2,
            )
        )
    else:
        report(sessions, by_model)


if __name__ == "__main__":
    main()