#!/usr/bin/env python3
"""Estimate token counts and cost for prompt files before you send them.

THIS IS AN APPROXIMATION. Every model family uses its own tokenizer, and the
default method here (characters divided by a ratio) can be off by 20% or more,
especially for code, non-English text, or heavy whitespace. Use it for budgeting
and spotting outliers, not for billing reconciliation. For exact counts use the
provider's own token counting endpoint or the usage numbers returned by the API.

Usage:
    python3 scripts/estimate_tokens.py FILE [FILE ...] [options]
    cat prompt.txt | python3 scripts/estimate_tokens.py -

Options:
    --chars-per-token R   Ratio for the default estimate (default 4.0, a common
                          rule of thumb for English prose).
    --tiktoken ENCODING   Count with the tiktoken library instead, e.g. o200k_base
                          or cl100k_base. These are OpenAI encodings; for other
                          vendors' models they are still only an approximation.
                          Requires: pip install tiktoken
    --output-tokens N     Expected output tokens per request (default 0).
    --requests N          Number of requests to project cost for (default 1).
    --input-price P       Your input price in USD per 1,000,000 tokens.
    --output-price P      Your output price in USD per 1,000,000 tokens.
                          Copy both from your provider's current pricing page;
                          no prices are built in because they change.
    --json                Print a JSON summary instead of a table.

Each FILE is treated as the input of one request. Totals add the files together.

Requirements: Python 3.8+. No environment variables. tiktoken is optional.

Exit codes:
    0  success
    2  usage error (unreadable file, bad option, tiktoken requested but missing)
"""
import argparse
import json
import math
import sys


def build_counter(args):
    if args.tiktoken:
        try:
            import tiktoken
        except ImportError:
            print("error: --tiktoken needs the tiktoken package (pip install tiktoken)", file=sys.stderr)
            sys.exit(2)
        try:
            encoding = tiktoken.get_encoding(args.tiktoken)
        except ValueError as exc:
            print("error: %s" % exc, file=sys.stderr)
            sys.exit(2)
        return "tiktoken:%s" % args.tiktoken, lambda text: len(encoding.encode(text, disallowed_special=()))
    if args.chars_per_token <= 0:
        print("error: --chars-per-token must be positive", file=sys.stderr)
        sys.exit(2)
    ratio = args.chars_per_token
    return "chars/%.2g" % ratio, lambda text: math.ceil(len(text) / ratio)


def read_input(path):
    if path == "-":
        return "<stdin>", sys.stdin.read()
    try:
        with open(path, encoding="utf-8", errors="replace") as handle:
            return path, handle.read()
    except OSError as exc:
        print("error: cannot read %s: %s" % (path, exc.strerror), file=sys.stderr)
        sys.exit(2)


def main():
    parser = argparse.ArgumentParser(description="Approximate token and cost estimator.")
    parser.add_argument("files", nargs="+")
    parser.add_argument("--chars-per-token", type=float, default=4.0)
    parser.add_argument("--tiktoken", metavar="ENCODING")
    parser.add_argument("--output-tokens", type=int, default=0)
    parser.add_argument("--requests", type=int, default=1)
    parser.add_argument("--input-price", type=float, help="USD per 1M input tokens")
    parser.add_argument("--output-price", type=float, help="USD per 1M output tokens")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()
    if args.output_tokens < 0 or args.requests < 1:
        print("error: --output-tokens must be >= 0 and --requests >= 1", file=sys.stderr)
        return 2

    method, count = build_counter(args)
    rows = []
    for path in args.files:
        name, text = read_input(path)
        rows.append({"file": name, "chars": len(text), "input_tokens": count(text)})

    input_per_request = sum(r["input_tokens"] for r in rows)
    total_input = input_per_request * args.requests
    total_output = args.output_tokens * args.requests
    summary = {
        "method": method,
        "approximate": True,
        "files": rows,
        "input_tokens_per_request": input_per_request,
        "output_tokens_per_request": args.output_tokens,
        "requests": args.requests,
        "total_input_tokens": total_input,
        "total_output_tokens": total_output,
    }
    if args.input_price is not None or args.output_price is not None:
        input_cost = total_input / 1_000_000 * (args.input_price or 0.0)
        output_cost = total_output / 1_000_000 * (args.output_price or 0.0)
        summary["estimated_cost_usd"] = round(input_cost + output_cost, 6)
        summary["estimated_input_cost_usd"] = round(input_cost, 6)
        summary["estimated_output_cost_usd"] = round(output_cost, 6)

    if args.json:
        print(json.dumps(summary, indent=2))
        return 0

    width = max([len(r["file"]) for r in rows] + [4])
    print("%-*s  %10s  %12s" % (width, "file", "chars", "~tokens"))
    for r in rows:
        print("%-*s  %10d  %12d" % (width, r["file"], r["chars"], r["input_tokens"]))
    print("\nmethod: %s (approximate)" % method)
    print("per request: ~%d input + %d output tokens" % (input_per_request, args.output_tokens))
    print("for %d request(s): ~%d input + %d output tokens" % (args.requests, total_input, total_output))
    if "estimated_cost_usd" in summary:
        print("estimated cost: ~$%.4f (input ~$%.4f, output ~$%.4f)" % (
            summary["estimated_cost_usd"], summary["estimated_input_cost_usd"], summary["estimated_output_cost_usd"]))
    else:
        print("cost: pass --input-price and --output-price (USD per 1M tokens) to estimate")
    return 0


if __name__ == "__main__":
    sys.exit(main())