#!/usr/bin/env python3
"""
根据 benchmark 结果计算 API 调用成本。

支持输入：
  - CSV（如 P800模型能力评测结果 - DS4-Flash-INT8-NO-Thinking-2.0-FULL.csv）
  - Excel（如 P800模型能力评测结果_统一格式_filled.xlsx）

计价公式：
  input_cost  = n_samples * input_tokens_mean / 1_000_000 * input_price * discount
  output_cost = n_samples * output_tokens_mean / 1_000_000 * output_price * discount

价格单位：元 / 百万 tokens（按国内 API 常见报价）。
"""

import argparse
import sys
from pathlib import Path

import pandas as pd


def load_input(path: str, sheet: str = None) -> pd.DataFrame:
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(path)

    if p.suffix.lower() in ('.xlsx', '.xls'):
        if sheet:
            return pd.read_excel(path, sheet_name=sheet)
        xl = pd.ExcelFile(path)
        print(f"可用 sheet: {xl.sheet_names}")
        return pd.read_excel(path, sheet_name=xl.sheet_names[0])
    else:
        return pd.read_csv(path)


def compute_cost(
    df: pd.DataFrame,
    input_price: float,
    output_price: float,
    discount: float,
    model_name: str,
) -> pd.DataFrame:
    df = df.copy()

    # 兼容中英文列名
    col_n_samples = next((c for c in df.columns if '样本' in c or c.lower() == 'n_samples'), None)
    col_input = next((c for c in df.columns if '输入tokens' in c or c.lower() == 'input_tokens_mean'), None)
    col_output = next((c for c in df.columns if '输出tokens' in c or c.lower() == 'output_tokens_mean'), None)
    col_total = next((c for c in df.columns if '累计总tokens' in c or c.lower() == 'total_tokens'), None)
    col_score = next((c for c in df.columns if '得分' in c or c.lower() == 'score'), None)
    col_bench = next((c for c in df.columns if 'benchmark' in c.lower()), None)

    if col_n_samples is None or col_input is None or col_output is None:
        raise ValueError("输入文件缺少必要列：总样本数 / 输入tokens_mean / 输出tokens_mean")

    price_in = input_price * discount / 1_000_000.0
    price_out = output_price * discount / 1_000_000.0

    # 数值化，空值填 0
    n = pd.to_numeric(df[col_n_samples], errors='coerce').fillna(0)
    in_tok = pd.to_numeric(df[col_input], errors='coerce').fillna(0)
    out_tok = pd.to_numeric(df[col_output], errors='coerce').fillna(0)

    df['input_tokens_total'] = (n * in_tok).astype('int64')
    df['output_tokens_total'] = (n * out_tok).astype('int64')
    df['input_cost_yuan'] = n * in_tok * price_in
    df['output_cost_yuan'] = n * out_tok * price_out
    df['total_cost_yuan'] = df['input_cost_yuan'] + df['output_cost_yuan']

    # 元/千 tokens 便于查看
    df['input_price_per_1k_tokens'] = price_in * 1000
    df['output_price_per_1k_tokens'] = price_out * 1000
    df['model_pricing'] = f"{model_name} (in={input_price}*{-discount}, out={output_price}*{discount})"

    return df


def save_output(df: pd.DataFrame, path: str):
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    if p.suffix.lower() in ('.xlsx', '.xls'):
        with pd.ExcelWriter(path, engine='openpyxl') as writer:
            df.to_excel(writer, sheet_name='cost', index=False)
    else:
        df.to_csv(path, index=False, encoding='utf-8-sig')
    print(f"结果已保存: {path}")


def main():
    parser = argparse.ArgumentParser(description="根据 benchmark 结果计算 API 成本")
    parser.add_argument('--input', '-i', required=True, help='输入 CSV/Excel 路径')
    parser.add_argument('--sheet', '-s', default=None, help='Excel sheet 名（默认第一个）')
    parser.add_argument('--input-price', type=float, default=8.0, help='输入单价（元/百万 tokens）')
    parser.add_argument('--output-price', type=float, default=28.0, help='输出单价（元/百万 tokens）')
    parser.add_argument('--discount', '-d', type=float, default=0.65, help='折扣率，如 0.65')
    parser.add_argument('--model-name', '-m', default='GLM-5.2', help='模型名称（仅用于标注）')
    parser.add_argument('--output', '-o', required=True, help='输出 CSV/Excel 路径')
    args = parser.parse_args()

    df = load_input(args.input, args.sheet)
    df_cost = compute_cost(
        df,
        input_price=args.input_price,
        output_price=args.output_price,
        discount=args.discount,
        model_name=args.model_name,
    )

    # 汇总行（如果原表没有 totals，追加一个）
    numeric_cols = ['input_cost_yuan', 'output_cost_yuan', 'total_cost_yuan']
    total_row = {c: df_cost[c].sum() for c in numeric_cols}
    total_row['Benchmark'] = 'TOTAL'
    total_row['model_pricing'] = args.model_name

    # 插入或更新 TOTAL 行
    is_total = df_cost['Benchmark'].astype(str).str.lower().isin(['total', '总计', '合计'])
    if is_total.any():
        for c in numeric_cols:
            df_cost.loc[is_total, c] = total_row[c]
    else:
        total_df = pd.DataFrame([total_row])
        df_cost = pd.concat([df_cost, total_df], ignore_index=True)

    save_output(df_cost, args.output)

    # 终端摘要
    summary = df_cost[['Benchmark', '总样本数', '输入tokens_mean', '输出tokens_mean',
                       'input_cost_yuan', 'output_cost_yuan', 'total_cost_yuan']].copy()
    print('\n=== 成本摘要（元）===')
    print(summary.to_string(index=False))
    print(f"\n总成本: {total_row['total_cost_yuan']:.4f} 元")


if __name__ == '__main__':
    main()