#!/usr/bin/env python3
import json
import sys


def walk(value):
    if isinstance(value, dict):
        yield value
        for nested in value.values():
            yield from walk(nested)
    elif isinstance(value, list):
        for nested in value:
            yield from walk(nested)


def number(value):
    return value if isinstance(value, (int, float)) and not isinstance(value, bool) else None


def normalize(candidate, model):
    if not isinstance(candidate, dict):
        return None
    input_tokens = candidate.get("inputTokens", candidate.get("input_tokens"))
    output_tokens = candidate.get("outputTokens", candidate.get("output_tokens"))
    cached_tokens = candidate.get("cachedTokens", candidate.get("cached_input_tokens", candidate.get("cache_read_input_tokens")))
    cost_usd = candidate.get("costUsd", candidate.get("cost_usd"))
    input_value, output_value = number(input_tokens), number(output_tokens)
    if input_value is None or output_value is None:
        return None
    usage = {"inputTokens": input_value, "outputTokens": output_value, "model": candidate.get("model") or model}
    cached_value, cost_value = number(cached_tokens), number(cost_usd)
    if cached_value is not None:
        usage["cachedTokens"] = cached_value
    if cost_value is not None:
        usage["costUsd"] = cost_value
    return usage


def main():
    raw_path, model = sys.argv[1], sys.argv[2]
    try:
        with open(raw_path, "r", encoding="utf-8") as handle:
            raw = handle.read()
    except OSError:
        raw = ""
    for line in raw.splitlines():
        line = line.strip()
        if not line.startswith("{"):
            continue
        try:
            parsed = json.loads(line)
        except ValueError:
            continue
        for candidate in walk(parsed):
            usage = normalize(candidate.get("usage"), model) if isinstance(candidate.get("usage"), dict) else None
            usage = usage or normalize(candidate, model)
            if usage is not None:
                print(json.dumps({"kind": "usage", "usage": usage}, separators=(",", ":")))
                return
    print(json.dumps({"kind": "usage", "usage": {"unknown": True, "model": model}}, separators=(",", ":")))


main()