import argparse
import glob
import json
import os


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--shard_glob", type=str, required=True)
    parser.add_argument("--output_path", type=str, required=True)
    return parser.parse_args()


def main():
    args = get_args()
    shard_paths = sorted(glob.glob(args.shard_glob))
    if not shard_paths:
        raise FileNotFoundError(f"No shard files matched: {args.shard_glob}")

    shards = [json.load(open(path, "r")) for path in shard_paths]
    model_name_red = shards[0]["model_name_red"]
    model_name_blue = shards[0]["model_name_blue"]
    file_name_red = shards[0]["file_name_red"]
    file_name_blue = shards[0]["file_name_blue"]

    evaluations = []
    for shard in shards:
        evaluations.extend(shard["evaluations"])
    evaluations.sort(key=lambda row: row.get("global_index", 10**9))

    win_h = tie_h = lose_h = nd_h = 0
    win_s = tie_s = lose_s = nd_s = 0
    for row in evaluations:
        rh, bh = row["red_score_helpfulness"], row["blue_score_helpfulness"]
        rs, bs = row["red_score_harmlessness"], row["blue_score_harmlessness"]
        if rh < 0 or bh < 0:
            nd_h += 1
        else:
            win_h += rh > bh
            tie_h += rh == bh
            lose_h += rh < bh
        if rs < 0 or bs < 0:
            nd_s += 1
        else:
            win_s += rs > bs
            tie_s += rs == bs
            lose_s += rs < bs

    denom_h = win_h + tie_h + lose_h
    denom_s = win_s + tie_s + lose_s
    result = {
        "file_name_red": file_name_red,
        "file_name_blue": file_name_blue,
        "model_name_red": model_name_red,
        "model_name_blue": model_name_blue,
        "win_helpfulness": win_h,
        "tie_helpfulness": tie_h,
        "lose_helpfulness": lose_h,
        "not_determined_helpfulness": nd_h,
        "win_rate_helpfulness": win_h / denom_h if denom_h else None,
        "win_halfTie_helpfulness": (win_h + 0.5 * tie_h) / denom_h if denom_h else None,
        "win_harmlessness": win_s,
        "tie_harmlessness": tie_s,
        "lose_harmlessness": lose_s,
        "not_determined_harmlessness": nd_s,
        "win_rate_harmlessness": win_s / denom_s if denom_s else None,
        "win_halfTie_harmlessness": (win_s + 0.5 * tie_s) / denom_s if denom_s else None,
        "evaluations": evaluations,
        "source_shards": shard_paths,
    }
    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    with open(args.output_path, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    main()