"""Calculate an ordinary text-token estimate locally; makes no network calls."""

import argparse
from decimal import Decimal, InvalidOperation


def nonnegative_decimal(value):
    try:
        number = Decimal(value)
    except InvalidOperation as exc:
        raise argparse.ArgumentTypeError("需要数字") from exc
    if not number.is_finite() or number < 0:
        raise argparse.ArgumentTypeError("需要有限的非负数字")
    return number


def nonnegative_int(value):
    try:
        number = int(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("Token 数量需要整数") from exc
    if number < 0:
        raise argparse.ArgumentTypeError("Token 数量不能为负数")
    return number


def positive_int(value):
    number = nonnegative_int(value)
    if number == 0:
        raise argparse.ArgumentTypeError("需要大于零的整数")
    return number


def estimate(input_tokens, output_tokens, input_price, output_price,
             price_unit_tokens=1000000, requests=1):
    unit = Decimal(price_unit_tokens)
    input_cost = Decimal(input_tokens) / unit * input_price
    output_cost = Decimal(output_tokens) / unit * output_price
    total = input_cost + output_cost
    return input_cost, output_cost, total, total * requests


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input-tokens", type=nonnegative_int, required=True)
    parser.add_argument("--output-tokens", type=nonnegative_int, required=True)
    parser.add_argument("--input-price", type=nonnegative_decimal, required=True,
                        help="输入价格 / 指定 Token 单位，与输出价格使用同一货币")
    parser.add_argument("--output-price", type=nonnegative_decimal, required=True,
                        help="输出价格 / 指定 Token 单位，与输入价格使用同一货币")
    parser.add_argument("--price-unit-tokens", type=positive_int, default=1000000,
                        help="价格对应的 Token 数量，默认 1000000；两种价格须同单位")
    parser.add_argument("--requests", type=positive_int, default=1,
                        help="重复同样用量的请求数，默认 1")
    args = parser.parse_args(argv)
    input_cost, output_cost, total, batch_total = estimate(
        args.input_tokens, args.output_tokens, args.input_price,
        args.output_price, args.price_unit_tokens, args.requests,
    )
    print(f"Input cost: {input_cost:f}")
    print(f"Output cost: {output_cost:f}")
    print(f"Total estimated cost: {total:f}")
    print(f"Requests: {args.requests}")
    print(f"Total estimated batch cost: {batch_total:f}")
    print("Estimate only. Excludes caching and other billing items; use actual usage and current prices.")


if __name__ == "__main__":
    main()