# Fixed PDF annotator
#
# Execute with:
#   python metabee/annotator.py --basepath data
#
# m.mieskolainen@imperial.ac.uk, 2025

import argparse
import json
import os

import fitz  # PyMuPDF
from termcolor import cprint


def convert_relative_to_absolute(page, rel_box):
    # Get page dimensions (width, height)
    width, height = page.rect.width, page.rect.height
    # Convert relative (0-1) coordinates to absolute values
    left = rel_box["l"] * width
    top = rel_box["t"] * height
    right = rel_box["r"] * width
    bottom = rel_box["b"] * height
    return fitz.Rect(left, top, right, bottom)


def annotate_pdf(pdf_path, merged_json_path, output_pdf, answers_json_path=None):
    """
    Annotate PDFs with bounding boxes.
    First, annotate all "question-answer" chunks (red boxes).
    Then, if an answers.json file is provided, gather field names for each chunk_id and annotate
    with blue boxes and a combined text label.
    """

    # Load merged JSON file.
    with open(merged_json_path, "r", encoding="utf-8") as f:
        merged = json.load(f)
    cprint(f"Loaded merged JSON: {merged_json_path}", "white")

    doc = fitz.open(pdf_path)

    # Annotate "question-answer" chunks with red rectangles.
    for chunk in merged["data"]["chunks"]:
        if chunk.get("chunk_type") == "question-answer":
            if "grounding" in chunk:
                for g in chunk["grounding"]:
                    page_num = g["page"]
                    if page_num < len(doc):
                        page = doc[page_num]
                        rect = convert_relative_to_absolute(page, g["box"])
                        # Draw a red rectangle with a width of 1.
                        page.draw_rect(rect, color=(1, 0, 0), width=1)

    # If answers.json is provided, process it.
    if answers_json_path and os.path.isfile(answers_json_path):
        with open(answers_json_path, "r", encoding="utf-8") as f:
            answers = json.load(f)
        cprint(f"Loaded answers JSON: {answers_json_path}", "cyan")

        # Build a dictionary mapping chunk_id to chunk for quick lookup.
        chunk_dict = {}
        for chunk in merged["data"]["chunks"]:
            cid = chunk.get("chunk_id")
            if cid:
                chunk_dict[cid] = chunk

        # Build a mapping from chunk_id to a set of field names.
        cid_to_fields = {}

        # Helper: Recursively extract chunk_ids from nested dictionaries.
        def extract_chunk_ids(d, current_label):
            if isinstance(d, dict):
                if "chunk_ids" in d:
                    for cid in d["chunk_ids"]:
                        cid_to_fields.setdefault(cid, set()).add(current_label)
                else:
                    for k, v in d.items():
                        extract_chunk_ids(v, k)
            elif isinstance(d, list):
                for item in d:
                    extract_chunk_ids(item, current_label)

        questions = answers.get("QUESTIONS", {})
        for question_key, question_value in questions.items():
            # For questions, either the field value directly has chunk_ids or search recursively.
            if isinstance(question_value, dict):
                for field_key, field_value in question_value.items():
                    if isinstance(field_value, dict) and "chunk_ids" in field_value:
                        for cid in field_value["chunk_ids"]:
                            cid_to_fields.setdefault(cid, set()).add(field_key)
                    else:
                        extract_chunk_ids(field_value, field_key)

        # Annotate each chunk from answers.json using the aggregated field names.
        for cid, fields in cid_to_fields.items():
            if cid in chunk_dict:
                chunk = chunk_dict[cid]
                if "grounding" in chunk:
                    for g in chunk["grounding"]:
                        page_num = g["page"]
                        if page_num < len(doc):
                            page = doc[page_num]
                            rect = convert_relative_to_absolute(page, g["box"])

                            # Draw a blue rectangle for answer-related chunks.
                            page.draw_rect(rect, color=(0, 0, 1), width=1)

                            # Insert a text annotation at the top left of the box.
                            field_text = ", ".join(sorted(fields))
                            annot_text = f"{cid}: ({field_text})"

                            shift = -5
                            page.insert_text(
                                (rect.x0, rect.y0 + shift), annot_text, fontname="helv", fontsize=5, color=(0, 0, 1)
                            )
            else:
                cprint(f"Warning: Chunk id {cid} not found in merged JSON", "yellow")
    else:
        cprint("No answers.json found - not processing answer-based annotations", "red")

    try:
        doc.save(output_pdf)
        cprint(f"Annotated PDF saved as: {output_pdf}", "green")
    except Exception as e:
        cprint(f"Error in saving PDF: {output_pdf}. Exception: {e}", "red")


def process_all_papers(base_papers_dir):
    """
    Process each paper folder (names that are digits) in sorted order.
    """
    paper_folders = sorted(
        [
            folder
            for folder in os.listdir(base_papers_dir)
            if os.path.isdir(os.path.join(base_papers_dir, folder)) and folder.isdigit()
        ]
    )

    for paper_folder in paper_folders:
        paper_path = os.path.join(base_papers_dir, paper_folder)
        pages_dir = os.path.join(paper_path, "pages")

        # Define file paths (adjust filenames as needed).
        original_pdf_path = os.path.join(paper_path, f"{paper_folder}_main.pdf")
        merged_json_path = os.path.join(pages_dir, "merged_v2.json")
        output_pdf = os.path.join(paper_path, f"{paper_folder}_main_annotated.pdf")
        answers_json_path = os.path.join(paper_path, "answers.json")

        # Check if necessary files exist.
        if not os.path.isfile(original_pdf_path):
            cprint(f"Paper {paper_folder}: Missing original PDF: {original_pdf_path}", "red")
            continue
        if not os.path.isfile(merged_json_path):
            cprint(f"Paper {paper_folder}: Missing merged JSON: {merged_json_path}", "red")
            continue

        # Annotate PDF (answers_json_path is optional).
        annotate_pdf(original_pdf_path, merged_json_path, output_pdf, answers_json_path)
        cprint(f"Paper {paper_folder}: Processing complete", "magenta")
        print()


def main():
    parser = argparse.ArgumentParser(description="Annotate PDFs using merged JSON and answers.json for papers.")
    parser.add_argument(
        "--basepath",
        type=str,
        default=os.getcwd(),
        help="Base path containing the 'papers' folder. Defaults to the current working directory.",
    )
    args = parser.parse_args()

    papers_dir = os.path.join(args.basepath, "papers")
    if not os.path.isdir(papers_dir):
        cprint(f"Error: 'papers' folder not found in {args.basepath}", "red")
        return
    process_all_papers(papers_dir)


if __name__ == "__main__":
    main()
