#!/usr/bin/env python3
import re
import json
import argparse

PAT_BATCH_START = re.compile(r"Batch Start \(id=(?P<bid>[^,]+), size=(?P<size>\d+)\).+req_ids=\[(?P<reqs>[^\]]*)\]")
PAT_META = re.compile(r"\[(INPAINT_META|REMOVEBG_META)\] (\{.*\})")


def parse_req_ids(s: str):
    s = s.strip()
    if not s:
        return []
    return [x.strip("' \"") for x in s.split(',') if x.strip()]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--log", required=True, help="path to main_size.log")
    args = ap.parse_args()

    batches = []
    metas = []
    with open(args.log, "r", encoding="utf-8") as f:
        for line in f:
            m1 = PAT_BATCH_START.search(line)
            if m1:
                reqs = parse_req_ids(m1.group("reqs"))
                batches.append({
                    "batch_id": m1.group("bid"),
                    "size": int(m1.group("size")),
                    "request_ids": reqs,
                    "line": line.strip(),
                })
            m2 = PAT_META.search(line)
            if m2:
                try:
                    metas.append(json.loads(m2.group(2)))
                except Exception:
                    pass

    # 매핑: request_id -> meta
    req_meta = {}
    for m in metas:
        rid = m.get("request_id")
        if rid:
            req_meta.setdefault(rid, []).append(m)

    # 결과 요약 출력
    print("Batches:", len(batches))
    for b in batches[-20:]:
        mapped = {rid: req_meta.get(rid, []) for rid in b["request_ids"]}
        print(json.dumps({
            "batch_id": b["batch_id"],
            "size": b["size"],
            "request_ids": b["request_ids"],
            "meta_found": {k: len(v) for k, v in mapped.items()},
        }, ensure_ascii=False))


if __name__ == "__main__":
    main()