import re
import argparse
from pathlib import Path
from typing import Dict, List, Tuple

SEP_RE = re.compile(r"^-{10,}\s*$")
IMG_RE = re.compile(r"^\s*IMAGE:\s*(.+?)\s*$")

# UPDATED: match "1." through "8."
Q_RE = re.compile(r"^\s*([1-8])\.\s*(.*)$")  # "1. ..." through "8. ..."

def parse_blocks(lines: List[str]) -> List[Tuple[str, Dict[int, str]]]:
    """
    Returns list of (image_path, answers_dict) where answers_dict maps qnum->answer_text.
    Counts a question as present even if answer_text is empty.
    """
    blocks = []
    cur_img = None
    cur_answers: Dict[int, List[str]] = {}
    cur_q = None

    def flush():
        nonlocal cur_img, cur_answers, cur_q
        if cur_img is None:
            return
        joined = {k: "\n".join(v).strip() for k, v in cur_answers.items()}
        blocks.append((cur_img, joined))
        cur_img = None
        cur_answers = {}
        cur_q = None

    for raw in lines:
        line = raw.rstrip("\n")

        # Separator means end of current block
        if SEP_RE.match(line):
            flush()
            continue

        m_img = IMG_RE.match(line)
        if m_img:
            # new IMAGE begins; flush previous if any
            flush()
            cur_img = m_img.group(1).strip()
            continue

        if cur_img is None:
            continue  # ignore anything before first IMAGE:

        m_q = Q_RE.match(line)
        if m_q:
            cur_q = int(m_q.group(1))
            cur_answers.setdefault(cur_q, [])
            remainder = m_q.group(2)
            cur_answers[cur_q].append(remainder.strip())
            continue

        # Continuation lines belong to the current question (if any)
        if cur_q is not None:
            cur_answers.setdefault(cur_q, [])
            cur_answers[cur_q].append(line.strip())

    # flush last block if file doesn't end with separator
    flush()
    return blocks

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--txt", required=True, help="Path to InternVL results txt file")
    ap.add_argument("--min_questions", type=int, default=7, help="Expected number of questions (default 8)")
    args = ap.parse_args()

    p = Path(args.txt)
    text = p.read_text(encoding="utf-8", errors="replace")
    blocks = parse_blocks(text.splitlines())

    print(f"Parsed {len(blocks)} IMAGE blocks from {p}")

    bad = []
    for img, answers in blocks:
        present = sorted(answers.keys())
        missing = [q for q in range(1, args.min_questions + 1) if q not in answers]
        if missing:
            bad.append((img, present, missing))

    if not bad:
        print("All images have answers for Q1–Q8.")
        return

    missing_image_names = []
    print(f"\nFound {len(bad)} images missing at least one answer:")
    for img, present, missing in bad:
        img_name = Path(img).name
        missing_image_names.append(img_name)
        print(f"- {img_name}  | present={present}  | missing={missing}")
    print("\nImages to regenerate (copy-paste list):")
    print("images_to_regenerate = [")
    for name in missing_image_names:
        print(f'    "{name}",')
    print("]")


if __name__ == "__main__":
    main()