#!/usr/bin/env python3
"""
为 P800模型能力评测结果 Excel 的每个 sheet 追加 API 成本列。
只处理包含 Benchmark + 总样本数 + 输入/输出 tokens_mean 的 sheet。
"""

import argparse
from pathlib import Path

import pandas as pd


def try_compute_cost(df: pd.DataFrame, input_price: float, output_price: float,
                     discount: float, model_name: str):
    col_n = next((c for c in df.columns if '样本' in c and '数' in c), None)
    col_in = next((c for c in df.columns if '输入tokens' in c), None)
    col_out = next((c for c in df.columns if '输出tokens' in c), None)
    col_bench = next((c for c in df.columns if 'benchmark' in c.lower()), None)

    if col_n is None or col_in is None or col_out is None or col_bench is None:
        return None

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

    n = pd.to_numeric(df[col_n], errors='coerce').fillna(0)
    in_tok = pd.to_numeric(df[col_in], errors='coerce').fillna(0)
    out_tok = pd.to_numeric(df[col_out], errors='coerce').fillna(0)

    result = df.copy()
    result['input_cost_元'] = n * in_tok * price_in
    result['output_cost_元'] = n * out_tok * price_out
    result['total_cost_元'] = result['input_cost_元'] + result['output_cost_元']
    result['price_model'] = model_name
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', '-i', required=True)
    parser.add_argument('--output', '-o', required=True)
    parser.add_argument('--input-price', type=float, default=8.0)
    parser.add_argument('--output-price', type=float, default=28.0)
    parser.add_argument('--discount', '-d', type=float, default=0.65)
    parser.add_argument('--model-name', '-m', default='GLM-5.2')
    args = parser.parse_args()

    xls = pd.ExcelFile(args.input)
    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    with pd.ExcelWriter(out_path, engine='openpyxl') as writer:
        for sheet in xls.sheet_names:
            df = pd.read_excel(args.input, sheet_name=sheet)
            new_df = try_compute_cost(
                df, args.input_price, args.output_price,
                args.discount, args.model_name,
            )
            if new_df is not None:
                # 汇总
                total_mask = new_df.iloc[:, 0].astype(str).str.lower().isin(['total', '总计', '合计'])
                if not total_mask.any():
                    total_row = {'price_model': args.model_name}
                    for c in new_df.columns:
                        if 'cost_元' in c:
                            total_row[c] = new_df[c].sum()
                    total_df = pd.DataFrame([total_row])
                    new_df = pd.concat([new_df, total_df], ignore_index=True)
                sheet_name = f"{sheet}_cost" if len(sheet) <= 25 else sheet[:25]
                new_df.to_excel(writer, sheet_name=sheet_name, index=False)
            else:
                # 原样复制
                df.to_excel(writer, sheet_name=sheet, index=False)

    print(f"已保存: {out_path}")


if __name__ == '__main__':
    main()