import argparse
import asyncio
import glob
import hashlib
import json
import logging
import os
import random
import re
import subprocess
import tempfile
import uuid
from collections import Counter, defaultdict
from typing import Dict, List

import pypdf
from anthropic import AsyncAnthropic
from bs4 import BeautifulSoup
from markdownify import SPACES, MarkdownConverter
from playwright.async_api import async_playwright
from syntok.segmenter import process
from tqdm import tqdm
from wordfreq import zipf_frequency

from olmocr.bench.tests import (
    BaselineTest,
    FootnoteTest,
    FormatTest,
    TableTest,
    TestType,
    TextOrderTest,
    TextPresenceTest,
    normalize_text,
    parse_html_tables,
)
from olmocr.data.renderpdf import (
    get_png_dimensions_from_base64,
    render_pdf_to_base64png,
)
from olmocr.filter.filter import Language, PdfFilter
from olmocr.synth.claude_client import (
    DEFAULT_MODEL_NAME,
    call_claude,
    claude_stream,
    extract_code_block,
)
from olmocr.synth.cutoff_detection import (
    RenderResult,
    _detect_cutoff_on_page,
    has_significant_cutoff,
)

# Global variables for tracking Claude API costs
total_input_tokens = 0
total_output_tokens = 0


def get_git_commit_hash():
    """Get the current git commit hash, if available."""
    try:
        result = subprocess.run(["git", "rev-parse", "HEAD"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
        return result.stdout.strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        # Git not available or not a git repository
        return None


# Unicode mappings for superscript characters
SUPERSCRIPT_MAP = {
    "0": "⁰",
    "1": "¹",
    "2": "²",
    "3": "³",
    "4": "⁴",
    "5": "⁵",
    "6": "⁶",
    "7": "⁷",
    "8": "⁸",
    "9": "⁹",
    "+": "⁺",
    "-": "⁻",
    "=": "⁼",
    "(": "⁽",
    ")": "⁾",
    "n": "ⁿ",
    "i": "ⁱ",
}

# Unicode mappings for subscript characters
SUBSCRIPT_MAP = {
    "0": "₀",
    "1": "₁",
    "2": "₂",
    "3": "₃",
    "4": "₄",
    "5": "₅",
    "6": "₆",
    "7": "₇",
    "8": "₈",
    "9": "₉",
    "+": "₊",
    "-": "₋",
    "=": "₌",
    "(": "₍",
    ")": "₎",
    "a": "ₐ",
    "e": "ₑ",
    "o": "ₒ",
    "x": "ₓ",
    "h": "ₕ",
    "k": "ₖ",
    "l": "ₗ",
    "m": "ₘ",
    "n": "ₙ",
    "p": "ₚ",
    "s": "ₛ",
    "t": "ₜ",
}


def convert_superscripts_subscripts(element):
    """
    Convert HTML superscript and subscript tags to Unicode equivalents.

    This function finds all <sup> and <sub> tags in the given element and
    replaces them with their Unicode character equivalents. Characters not
    in the mapping are left unchanged.

    Args:
        element: A BeautifulSoup element to process

    Returns:
        The element with sup/sub tags converted to Unicode
    """
    if not element:
        return element

    # Process all superscript tags
    for sup in element.find_all("sup"):
        sup_text = sup.get_text()
        unicode_text = "".join(SUPERSCRIPT_MAP.get(char, char) for char in sup_text)
        sup.replace_with(unicode_text)

    # Process all subscript tags
    for sub in element.find_all("sub"):
        sub_text = sub.get_text()
        unicode_text = "".join(SUBSCRIPT_MAP.get(char, char) for char in sub_text)
        sub.replace_with(unicode_text)

    return element


def download_s3_pdf(path, local_path):
    """Download a PDF from S3 or copy from local path."""
    os.makedirs(os.path.dirname(local_path), exist_ok=True)

    # Check if it's a local path
    if os.path.exists(path):
        # It's a local file, just copy it
        import shutil

        try:
            shutil.copy2(path, local_path)
            return True
        except Exception as e:
            print(f"Failed to copy local file {path}: {e}")
            return False
    elif path.startswith("s3://"):
        # It's an S3 path, download it
        result = subprocess.run(["aws", "s3", "cp", path, local_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        return result.returncode == 0
    else:
        # Assume it's a relative local path that doesn't exist yet
        print(f"Path not found and doesn't appear to be S3: {path}")
        return False


def cleanup_headers_footers_soup(soup):
    # Remove headers completely
    for header in soup.find_all("header"):
        header.decompose()

    # For footers: remove direct text but keep footnote elements
    for footer in soup.find_all("footer"):
        # First, preserve all footnote elements (div, span, p with class="footnote")
        footnote_elements = []
        for tag_type in ["div", "span", "p"]:
            footnote_elements.extend(footer.find_all(tag_type, class_="footnote"))

        # Extract and temporarily store footnote elements
        preserved_elements = []
        for fn_element in footnote_elements:
            # Extract the element from its current position
            fn_element.extract()
            preserved_elements.append(fn_element)

        # Clear all content from the footer
        footer.clear()

        # Re-add only the footnote elements back to the footer
        for fn_element in preserved_elements:
            footer.append(fn_element)

    # Remove any divs or spans with class "line-number"
    for element in soup.find_all(["div", "span"], class_="line-number"):
        element.extract()

    # Remove any div or span watermarks
    for element in soup.find_all(["div", "span"], class_="watermark"):
        element.extract()


class PreserveTablesConverter(MarkdownConverter):
    """
    Custom MarkdownConverter that preserves HTML tables unchanged
    and preserves sup/sub tags as HTML
    """

    def convert_table(self, el, text, parent_tags):
        # Get the outer HTML of the table element
        # BeautifulSoup's prettify or str() should give us the full HTML
        from bs4 import BeautifulSoup

        # Create a temporary soup with just this element to get its HTML
        temp_soup = BeautifulSoup(str(el), "html.parser")
        return str(temp_soup.table) if temp_soup.table else str(el)

    def convert_sup(self, el, text, parent_tags):
        # Always preserve sup tags as HTML
        return f"<sup>{el.get_text()}</sup>"

    def convert_sub(self, el, text, parent_tags):
        # Always preserve sub tags as HTML
        return f"<sub>{el.get_text()}</sub>"


def extract_html_metadata(html_content):
    """Extract metadata from HTML content for FrontMatter."""
    soup = BeautifulSoup(html_content, "html.parser")

    # Extract language from html tag
    html_tag = soup.find("html")
    language = "en"  # default
    if html_tag and html_tag.get("lang"):
        language = str(html_tag.get("lang"))
        # Convert pt-BR to pt for now
        if len(language) == 5 and language[2] == "-":
            language = language[:2]

    # Calculate content statistics
    body = soup.find("body")
    if not body:
        body = soup

    # First, create a version without headers and footers for all calculations
    main_content_soup = BeautifulSoup(str(body), "html.parser")
    # Remove headers and footers from main content
    for element in main_content_soup.find_all(["header", "footer"]):
        element.decompose()

    # Get text content length (excluding tables and images)
    text_soup = BeautifulSoup(str(main_content_soup), "html.parser")
    # Remove tables
    for element in text_soup.find_all("table"):
        element.decompose()
    # Remove images (div.image)
    for element in text_soup.find_all("div", class_="image"):
        element.decompose()
    text_content = text_soup.get_text().strip()
    text_length = len(text_content)

    # Count table content (from main content, excluding headers/footers)
    tables = main_content_soup.find_all("table")
    table_text_length = 0
    for table in tables:
        table_text_length += len(table.get_text().strip())

    # Count images (div.image elements) (from main content, excluding headers/footers)
    images = main_content_soup.find_all("div", class_="image")
    # Rough estimate: each image takes up about 500 characters worth of "space"
    image_content_estimate = len(images) * 500

    # Calculate total content "length"
    total_content_length = text_length + table_text_length + image_content_estimate

    # Determine if mostly tables or images
    is_table = False
    is_diagram = False

    if total_content_length > 0:
        table_ratio = table_text_length / total_content_length
        image_ratio = image_content_estimate / total_content_length

        is_table = table_ratio > 0.5
        is_diagram = image_ratio > 0.5

    return {"primary_language": language, "is_rotation_valid": True, "rotation_correction": 0, "is_table": is_table, "is_diagram": is_diagram}


def html_to_markdown_with_frontmatter(html_content):
    """Convert HTML to markdown with FrontMatter metadata."""
    # Extract metadata
    metadata = extract_html_metadata(html_content)

    # Parse HTML and extract only body content for markdown conversion
    soup = BeautifulSoup(html_content, "html.parser")
    body = soup.find("body")

    # If no body tag, use the whole soup as fallback
    if body:
        # Create a new soup with just the body content
        body_soup = BeautifulSoup(str(body), "html.parser")
    else:
        body_soup = soup

    # First, remove all header and footer elements from the body
    cleanup_headers_footers_soup(body_soup)

    # Also remove divs with page-header or page-footer classes (in case they weren't converted to header/footer tags)
    for div in body_soup.find_all("div", class_="page-header"):
        div.decompose()
    for div in body_soup.find_all("div", class_="page-footer"):
        div.decompose()

    # Handle image placeholders - replace div.image with actual img tags for proper markdown conversion
    for img_div in body_soup.find_all("div", class_="image"):
        alt_text = "Image Placeholder"  # For now, in the render it's all just a placeholder
        # Create an img tag with placeholder src and appropriate alt text
        img_tag = body_soup.new_tag("img", src="page.png", alt=alt_text)
        img_div.replace_with(img_tag)

    # Handle SVG pictures in a similar way, just replace it as an image tag
    for svg_tag in body_soup.find_all("svg"):
        alt_text = "Graphic Placeholder"
        img_tag = body_soup.new_tag("img", src="page.png", alt=alt_text)
        svg_tag.replace_with(img_tag)

    # Get the modified HTML (only body content)
    modified_html = str(body_soup)

    # Create custom converter instance
    converter = PreserveTablesConverter(
        heading_style="ATX",  # Use # style headings
        bullets="-",  # Use - for unordered lists
        strip=["a"],  # Remove links but keep text
        newline_style=SPACES,  # Use backslash for line breaks
        code_language="",  # Don't add language to code blocks
        escape_asterisks=False,  # Don't escape asterisks
        escape_underscores=False,  # Don't escape underscores
    )

    # Convert to markdown
    markdown = converter.convert(modified_html)

    # Clean up excessive newlines
    while "\n\n\n" in markdown:
        markdown = markdown.replace("\n\n\n", "\n\n")

    # Strip and clean up markdown content
    markdown_content = markdown.strip()

    # Remove leading or trailing --- if present
    while markdown_content.startswith("---"):
        markdown_content = markdown_content[3:].strip()
    while markdown_content.endswith("---"):
        markdown_content = markdown_content[:-3].strip()

    # Create FrontMatter
    frontmatter = f"""---
primary_language: {metadata['primary_language']}
is_rotation_valid: {metadata['is_rotation_valid']}
rotation_correction: {metadata['rotation_correction']}
is_table: {metadata['is_table']}
is_diagram: {metadata['is_diagram']}
---"""

    # Combine FrontMatter with markdown content
    if markdown_content:
        return f"{frontmatter}\n{markdown_content}"
    else:
        return frontmatter


async def generate_html_from_image(client, image_base64):
    """Call Claude API to generate HTML from an image using a multi-step prompting strategy."""
    global total_input_tokens, total_output_tokens
    png_width, png_height = get_png_dimensions_from_base64(image_base64)

    try:
        # Step 0: Check that the orientation of the original document is right-side-up. If not, we will
        # skip this page, to keep the code simple
        orientation_response = await call_claude(
            client,
            model=DEFAULT_MODEL_NAME,
            max_tokens=1000,
            temperature=0,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": image_base64}},
                        {
                            "type": "text",
                            "text": "Please analyze this document image and determine its orientation.\n\n"
                            "Is this document right-side-up (correctly oriented), or is it rotated?\n\n"
                            "Make your decision based on the main document contents that takes up most of the page area.\n\n"
                            "Respond with ONLY one of the following:\n"
                            "- RIGHT_SIDE_UP: The document is correctly oriented and readable\n"
                            "- ROTATED_90: The document is rotated 90 degrees clockwise\n"
                            "- ROTATED_180: The document is upside down (rotated 180 degrees)\n"
                            "- ROTATED_270: The document is rotated 270 degrees clockwise (90 degrees counter-clockwise)\n"
                            "- UNCLEAR: Cannot determine orientation (e.g., blank page, purely graphical content)\n\n"
                            "Important: Only respond with one of these exact terms, nothing else.",
                        },
                    ],
                }
            ],
        )

        # Extract orientation from response
        orientation_text = ""
        for content in orientation_response.content:
            if content.type == "text":
                orientation_text += content.text.strip()

        # Track token usage from orientation check
        if hasattr(orientation_response, "usage"):
            total_input_tokens += orientation_response.usage.input_tokens
            total_output_tokens += orientation_response.usage.output_tokens

        # Check orientation result
        if "RIGHT_SIDE_UP" not in orientation_text:
            print(f"Skipping page due to orientation: {orientation_text}")
            return None

        # Step 1: Initial analysis and column detection
        analysis_response = await call_claude(
            client,
            model=DEFAULT_MODEL_NAME,
            max_tokens=20000,
            temperature=0.1,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": image_base64}},
                        {
                            "type": "text",
                            "text": "Analyze this document and provide a detailed assessment of its structure. Focus specifically on:\n"
                            "1. How many columns does the document have? Is it single-column, two-column, three-column, or a mixed layout?\n"
                            "2. What are the main sections and content types (headings, paragraphs, lists, tables, images, etc.)?\n"
                            "3. Does it have headers, footers, page numbers, or other special elements?\n"
                            "4. Is there any complex formatting that would be challenging to reproduce in HTML?\n\n"
                            "Please be very precise about the number of columns and how they're arranged.",
                        },
                    ],
                }
            ],
        )

        # Check if response was complete
        if hasattr(analysis_response, "stop_reason") and analysis_response.stop_reason != "end_turn":
            print(f"Warning: Analysis response incomplete (stop_reason: {analysis_response.stop_reason})")
            return None

        analysis_text = ""
        for content in analysis_response.content:
            if content.type == "text":
                analysis_text += content.text

        # Track token usage from first API call
        if hasattr(analysis_response, "usage"):
            total_input_tokens += analysis_response.usage.input_tokens
            total_output_tokens += analysis_response.usage.output_tokens

        # Step 2: Initial HTML generation with detailed layout instructions
        initial_response = await call_claude(
            client,
            model=DEFAULT_MODEL_NAME,
            max_tokens=20000,
            temperature=0.2,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": image_base64}},
                        {
                            "type": "text",
                            "text": "Render this document as clean, semantic HTML. Here's my analysis of the document structure:\n\n"
                            f"{analysis_text}\n\n"
                            "Important requirements:\n"
                            "1. Use appropriate HTML tags for elements like headings, paragraphs, lists, tables, etc.\n"
                            "2. Use the <header> and <footer> tags to represent content at the top/bottom which would not normally be part of the main content, such as page numbers, etc.\n"
                            "3. Use a placeholder <div> tag with class 'image' which will render as a grey box with black outline to make sure images have their original size, shape, and position on the page. Include an alt-text of the original image as a 'data-description' attribute on the tag. Include 'data-x', 'data-y', 'data-width', 'data-height' attributes which specify where the image was found in the original document.\n"
                            "4. Render any math equations and Latex inline using either \\[ \\] or \\( \\) delimeters.\n"
                            "5. Render any subscripts and superscripts in <sub> and <sup> tags, not in Unicode characters.\n"
                            "6. CRITICAL: If the document has a multi-column layout, you MUST preserve the exact same number of columns in your HTML. Use CSS flexbox or grid to create the columns.\n"
                            "7. Focus on creating valid, accessible HTML that preserves the appearance and formatting of the original page as closely as possible.\n"
                            f"8. The webpage will be viewed with a fixed viewport size of {png_width} pixels wide by {png_height} pixels tall.\n"
                            "9. For multi-column layouts, use explicit CSS. The most important aspect is preserving the column structure of the original document - this is critical.\n\n"
                            "Enclose your HTML in a ```html code block.",
                        },
                    ],
                }
            ],
        )

        # Check if response was complete
        if hasattr(initial_response, "stop_reason") and initial_response.stop_reason != "end_turn":
            print(f"Warning: Initial HTML response incomplete (stop_reason: {initial_response.stop_reason})")
            return None

        # Extract initial HTML
        initial_html_text = ""
        for content in initial_response.content:
            if content.type == "text":
                initial_html_text += content.text

        # Track token usage from second API call
        if hasattr(initial_response, "usage"):
            total_input_tokens += initial_response.usage.input_tokens
            total_output_tokens += initial_response.usage.output_tokens

        initial_html = extract_code_block(initial_html_text)
        if not initial_html:
            print("Warning: No HTML code block found in initial response")
            return None

        # Step 3: Render the initial HTML to PDF and then back to PNG for comparison
        # Create a temporary PDF file
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp_pdf:
            tmp_pdf_path = tmp_pdf.name

        try:
            # Render HTML to PDF using existing function
            render_result = await render_pdf_with_playwright(initial_html, tmp_pdf_path, png_width, png_height)

            if not render_result.success:
                print("Warning: Failed to render initial HTML to PDF for refinement")
                return None

            # Convert PDF back to PNG
            rendered_image_base64 = render_pdf_to_base64png(tmp_pdf_path, 1, max(png_width, png_height))

            if not rendered_image_base64:
                print("Warning: Failed to convert rendered PDF to PNG for refinement")
                return None

            # We are going to add some stuff to the prompt conditioned on if tables need to be corrected or not
            extra_table_fixing_instructions = ""

            # Check the tables, if they are non-rectangular, we can apply one more correction pass on them
            table_data = parse_html_tables(initial_html)

            if any(not table.is_rectangular for table in table_data):
                extra_table_fixing_instructions = (
                    "Important: I've noticed that in the HTML table code, some of the columns/rows are not aligned right. "
                    "Please work extra hard to make sure the table columns are correctly lined up as in the original document. "
                    "You can add HTML comments as you output the table to help keep track of the current row and column if needed.\n"
                )

            # Step 4: Refinement - Show both images to Claude and ask for corrections
            refinement_response = await claude_stream(
                client,
                model=DEFAULT_MODEL_NAME,
                max_tokens=40000,
                temperature=1.0,
                thinking={"type": "enabled", "budget_tokens": 12000},
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": "I'm going to show you two images:\n1. The original document\n2. How the HTML I generated renders\n\nPlease compare them carefully and provide a revised version of the HTML that better matches the original.",
                            },
                            {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": image_base64}},
                            {"type": "text", "text": "Above is the ORIGINAL document."},
                            {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": rendered_image_base64}},
                            {"type": "text", "text": "Above is how my HTML currently renders."},
                            {
                                "type": "text",
                                "text": f"Here is the current HTML code:\n\n```html\n{initial_html}\n```\n\n"
                                "Please analyze the differences between the original document and the rendered version. Focus on:\n"
                                "1. Layout issues - are columns preserved correctly?\n"
                                "2. Positioning - are elements in the right place?\n"
                                "3. Spacing - are margins, padding, and spacing between elements correct?\n"
                                "4. Occlusion - is any important content hidden or overlapping?\n"
                                "5. Text formatting - are fonts, sizes, and styles appropriate?\n"
                                "6. Math - are any Latex math equations delimited by either \\[ \\] or \\( \\) delimeters?\n"
                                "7. Tables - are the headers on tables are aligned with the correct corresponding columns?\n"
                                f"{extra_table_fixing_instructions}"
                                f"The webpage will be viewed at {png_width}x{png_height} pixels.\n\n"
                                "Provide a REVISED version of the HTML that corrects any issues you identified. "
                                "Make sure all important elements are visible and the layout matches the original as closely as possible.\n"
                                "Output the complete revised HTML in a ```html code block.",
                            },
                        ],
                    }
                ],
            )

            # Check if refinement response was complete
            if hasattr(refinement_response, "stop_reason") and refinement_response.stop_reason != "end_turn":
                print(f"Warning: Refinement response incomplete (stop_reason: {refinement_response.stop_reason})")
                # Return initial HTML as fallback since it was complete
                return initial_html

            # Extract refined HTML
            refined_html_text = ""
            for content in refinement_response.content:
                if content.type == "text":
                    refined_html_text += content.text

            # Track token usage from refinement API call
            if hasattr(refinement_response, "usage"):
                total_input_tokens += refinement_response.usage.input_tokens
                total_output_tokens += refinement_response.usage.output_tokens

            refined_html = extract_code_block(refined_html_text)
            final_html = refined_html if refined_html else initial_html

            # Check the tables, if they are non-rectangular, we can apply one more correction pass on them
            table_data = parse_html_tables(final_html)

            if any(not table.is_rectangular for table in table_data):
                print("Table not rectangular, aborting")
                return None

            # Return refined HTML if available, otherwise return initial HTML
            return final_html

        finally:
            # Clean up temporary PDF file
            if os.path.exists(tmp_pdf_path):
                os.remove(tmp_pdf_path)

    except Exception as e:
        print(f"Error calling Claude API: {e}")
        return None


def extract_page_from_pdf(input_path, output_path, page_num):
    """
    Extract a specific page from a PDF and save it as a new PDF.

    Args:
        input_path: Path to the input PDF
        output_path: Path to save the extracted page
        page_num: The page number to extract (1-indexed, converted to 0-indexed for pypdf)

    Returns:
        bool: True if extraction was successful, False otherwise
    """
    try:
        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # Read the input PDF
        reader = pypdf.PdfReader(input_path)

        # Convert to 0-indexed for pypdf
        zero_idx_page = page_num - 1

        # Check if page number is valid
        if zero_idx_page >= len(reader.pages) or zero_idx_page < 0:
            print(f"Page number {page_num} out of range for {input_path} with {len(reader.pages)} pages")
            return False

        # Create a new PDF with just the selected page
        writer = pypdf.PdfWriter()
        writer.add_page(reader.pages[zero_idx_page])

        # Write the output PDF
        with open(output_path, "wb") as output_file:
            writer.write(output_file)

        return True
    except Exception as e:
        print(f"Error extracting page {page_num} from {input_path}: {str(e)}")
        return False


async def _load_katex_on_page(page):
    """Load KaTeX CSS/JS and run auto-render on an already-loaded Playwright page."""
    katex_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "bench", "katex")
    katex_css_path = os.path.join(katex_dir, "katex.min.css")
    katex_js_path = os.path.join(katex_dir, "katex.min.js")
    katex_autorender_js_path = os.path.join(katex_dir, "auto-render.min.js")

    await page.add_style_tag(path=katex_css_path)
    await page.add_script_tag(path=katex_js_path)
    await page.add_script_tag(path=katex_autorender_js_path)

    await page.evaluate("""
        renderMathInElement(document.body, {
            delimiters: [
                {left: '\\\\(', right: '\\\\)', display: false},
                {left: '\\\\[', right: '\\\\]', display: true}
            ],
            throwOnError: false
        });
    """)


async def render_pdf_with_playwright(html_content, output_pdf_path, png_width, png_height):
    """
    Render HTML content using Playwright and save it as PDF.
    First checks for text cutoff (overflow clipping), then tries different
    scale factors if needed to ensure the output is exactly one page.

    Args:
        html_content: HTML content to render
        output_pdf_path: Path to save the rendered PDF
        png_width: Width of the viewport
        png_height: Height of the viewport

    Returns:
        RenderResult with success, scale_used, cutoff info
    """
    scale_factors = [1.0, 0.9, 0.8, 0.7, 0.6, 0.5]

    # Determine page format based on PNG dimensions
    aspect_ratio = png_width / png_height

    pdf_options = {
        "path": output_pdf_path,
        "print_background": True,
    }

    if 0.73 <= aspect_ratio <= 0.81:  # Letter Portrait
        pdf_options["width"] = "8.5in"
        pdf_options["height"] = "11in"
    elif 1.23 <= aspect_ratio <= 1.35:  # Letter Landscape
        pdf_options["width"] = "11in"
        pdf_options["height"] = "8.5in"
    elif 0.67 <= aspect_ratio <= 0.73:  # A4 Portrait
        pdf_options["width"] = "210mm"
        pdf_options["height"] = "297mm"
    elif 1.36 <= aspect_ratio <= 1.47:  # A4 Landscape
        pdf_options["width"] = "297mm"
        pdf_options["height"] = "210mm"

    async with async_playwright() as p:
        browser = await p.chromium.launch()

        try:
            # Phase 1: Cutoff detection at original viewport size
            check_page = await browser.new_page(viewport={"width": png_width, "height": png_height})
            await check_page.set_content(html_content, wait_until="load")
            await _load_katex_on_page(check_page)

            cutoff_elements = await _detect_cutoff_on_page(check_page, 0.9)
            await check_page.close()

            if has_significant_cutoff(cutoff_elements):
                await browser.close()
                return RenderResult(
                    success=False,
                    cutoff_elements=cutoff_elements,
                    has_cutoff=True,
                )

            # Phase 2: Scale loop — render PDF, check page count
            for scale in scale_factors:
                try:
                    page = await browser.new_page(
                        viewport={
                            "width": int(png_width * scale),
                            "height": int(png_height * scale),
                        }
                    )
                    await page.set_content(html_content)
                    await _load_katex_on_page(page)

                    pdf_options["scale"] = scale
                    await page.pdf(**pdf_options)
                    await page.close()

                    try:
                        reader = pypdf.PdfReader(output_pdf_path)
                        if len(reader.pages) == 1:
                            print(f"Successfully rendered as a single page PDF with scale factor {scale}")
                            await browser.close()
                            return RenderResult(success=True, scale_used=scale)
                        else:
                            print(f"PDF has {len(reader.pages)} pages with scale factor {scale}, trying a smaller scale...")
                    except Exception as pdf_check_error:
                        print(f"Error checking PDF page count: {pdf_check_error}")
                        await browser.close()
                        return RenderResult(success=False)

                except Exception as e:
                    print(f"Error rendering PDF with Playwright at scale {scale}: {str(e)}")

        except Exception as e:
            print(f"Error during render_pdf_with_playwright: {str(e)}")
            await browser.close()
            return RenderResult(success=False)

        await browser.close()

    print("Failed to render PDF as a single page with any scale factor")
    return RenderResult(success=False)


def generate_tests_from_html(html_content: str, pdf_id: str, page_num: int, random_gen: random.Random, verbose_table_testing: bool = False) -> List[Dict]:
    """
    Generate tests from HTML content parsed from the PDF.

    Args:
        html_content: The HTML content of the page
        pdf_id: The unique identifier for the PDF
        page_num: The page number
        verbose_table_testing: Whether to print table test verification details

    Returns:
        A list of test dictionaries that can be saved as JSONL
    """

    # Use the module-level conversion function

    tests = []
    pdf_filename = f"{pdf_id}_page{page_num}.pdf"
    soup = BeautifulSoup(html_content, "html.parser")

    # Remove any divs or spans with class "line-number"
    for element in soup.find_all(["div", "span"], class_="line-number"):
        element.extract()

    # Rewrite any page-header and page-footer divs to be normalized to headers
    # Convert div.page-footer to footer in one line
    for div in soup.find_all("div", class_="page-header"):
        div.name = "header"

    for div in soup.find_all("div", class_="page-footer"):
        div.name = "footer"

    # Remove elements in the body that appear before the header or after the footer
    body = soup.find("body")
    if body:
        header = soup.find("header")
        footer = soup.find("footer")

        if header:
            # Remove elements before the header
            current = body.contents[0]
            while current and current != header:
                next_elem = current.next_sibling
                current.extract()
                current = next_elem

        if footer:
            # Remove elements after the footer
            current = footer.next_sibling
            while current:
                next_elem = current.next_sibling
                current.extract()
                current = next_elem

    # Step 1: Process headers, footers, and page numbers for TextAbsenceTests
    headers = soup.find_all("header")
    footers = soup.find_all("footer")
    page_numbers = soup.find_all("div", class_="page-number")

    # Function to create absence tests from text elements
    def create_absence_tests_from_elements(parent_element, element_type):
        mini_soup = BeautifulSoup(str(parent_element), "html.parser")

        # Convert superscripts and subscripts in the mini soup
        convert_superscripts_subscripts(mini_soup)

        # Remove headers, footers, and tables from the main_soup
        for element in mini_soup.find_all(["h1", "h2"]):
            element.extract()

        # Find all text-containing leaf elements within the parent
        text_elements = []

        # Get all target elements
        target_tags = mini_soup.find_all(["span", "div", "p", "h3", "h4", "h5", "h6"])

        # Filter to only include leaf nodes (elements that don't contain other target elements)
        for tag in target_tags:
            # Check if this element has no children from our target tags
            is_leaf = not tag.find(["span", "div", "p", "h3", "h4", "h5", "h6"])

            if is_leaf:
                text = tag.get_text().strip()
                if text:
                    text_elements.append(text)

        # If no elements found, use the parent's text as a fallback, but only if
        if not text_elements:
            parent_text = mini_soup.get_text().strip()
            if parent_text:
                text_elements.append(parent_text)

        # Create tests for each text element
        for text in text_elements:
            if "\n" in text:
                text = text.split("\n")[0]

            if len(text) > 3 or len([c for c in text if c.isdigit()]):  # Only create tests for meaningful text
                tests.append(
                    {
                        "pdf": pdf_filename,
                        "page": 1,
                        "id": f"{pdf_id}_{element_type}_{uuid.uuid4().hex[:8]}",
                        "type": TestType.ABSENT.value,
                        "text": text,
                        "max_diffs": round(len(text) * 0.05),
                    }
                )

    # Create TextAbsenceTests for headers
    for header in headers:
        create_absence_tests_from_elements(header, "header")

    # Create TextAbsenceTests for footers
    for footer in footers:
        create_absence_tests_from_elements(footer, "footer")

    # Create TextAbsenceTests for page numbers
    for page_number in page_numbers:
        # Convert any superscripts/subscripts in the page number
        page_number_soup = BeautifulSoup(str(page_number), "html.parser")
        convert_superscripts_subscripts(page_number_soup)
        page_number_text = page_number_soup.get_text().strip()

        if page_number_text:
            tests.append(
                {
                    "pdf": pdf_filename,
                    "page": 1,
                    "id": f"{pdf_id}_page_number_{uuid.uuid4().hex[:8]}",
                    "type": TestType.ABSENT.value,
                    "text": page_number_text,
                    "max_diffs": 0,
                }
            )

    # Step 2: Generate tests from tables using parse_html_tables
    # Convert superscripts and subscripts to Unicode equivalents in tables
    table_soup = BeautifulSoup(html_content, "html.parser")

    # Convert superscripts and subscripts in the table HTML
    convert_superscripts_subscripts(table_soup)
    html_content_with_unicode = str(table_soup)

    table_data_list = parse_html_tables(html_content_with_unicode)

    for table_idx, table_data in enumerate(table_data_list):
        # Get the table data as a numpy array
        table_tests = []

        # Skip tables that are too small
        if max(x[0] for x in table_data.cell_text.keys()) < 2:
            continue
        if max(x[1] for x in table_data.cell_text.keys()) < 2:
            continue

        all_known_cells = list(table_data.cell_text.items())
        random_gen.shuffle(all_known_cells)

        for rowcol, cell_content in all_known_cells:
            cell_content = normalize_text(cell_content)

            if len(cell_content) == 0:
                continue

            # Create a TableTest with relevant relationships
            test_data = {
                "pdf": pdf_filename,
                "page": 1,
                "id": f"{pdf_id}_table{table_idx}_{uuid.uuid4().hex[:8]}",
                "type": TestType.TABLE.value,
                "cell": cell_content,
                "max_diffs": 0,
                "ignore_markdown_tables": True,
            }

            if rowcol in table_data.up_relations and len(table_data.up_relations[rowcol]) > 0:
                relation = random_gen.choice(list(table_data.up_relations[rowcol]))
                if len(table_data.cell_text[relation].strip()) > 0 and "\n" not in table_data.cell_text[relation]:
                    test_data["up"] = normalize_text(table_data.cell_text[relation])

            if rowcol in table_data.down_relations and len(table_data.down_relations[rowcol]) > 0:
                relation = random_gen.choice(list(table_data.down_relations[rowcol]))
                if len(table_data.cell_text[relation].strip()) > 0 and "\n" not in table_data.cell_text[relation]:
                    test_data["down"] = normalize_text(table_data.cell_text[relation])

            if rowcol in table_data.left_relations and len(table_data.left_relations[rowcol]) > 0:
                relation = random_gen.choice(list(table_data.left_relations[rowcol]))
                if len(table_data.cell_text[relation].strip()) > 0 and "\n" not in table_data.cell_text[relation]:
                    test_data["left"] = normalize_text(table_data.cell_text[relation])

            if rowcol in table_data.right_relations and len(table_data.right_relations[rowcol]) > 0:
                relation = random_gen.choice(list(table_data.right_relations[rowcol]))
                if len(table_data.cell_text[relation].strip()) > 0 and "\n" not in table_data.cell_text[relation]:
                    test_data["right"] = normalize_text(table_data.cell_text[relation])

            if len(table_data.left_heading_relations(*rowcol)) > 0:
                relation = random_gen.choice(list(table_data.left_heading_relations(*rowcol)))
                if len(table_data.cell_text[relation].strip()) > 0 and "\n" not in table_data.cell_text[relation]:
                    test_data["left_heading"] = normalize_text(table_data.cell_text[relation])

            if len(table_data.top_heading_relations(*rowcol)) > 0:
                relation = random_gen.choice(list(table_data.top_heading_relations(*rowcol)))
                if len(table_data.cell_text[relation].strip()) > 0 and "\n" not in table_data.cell_text[relation]:
                    test_data["top_heading"] = normalize_text(table_data.cell_text[relation])

            # Only add the test if we have at least one relation
            if any(x in test_data for x in ["up", "down", "left", "right", "top_heading", "left_heading"]):
                # Verify that the test passes with the current table HTML
                # Create the actual test object
                test_obj = TableTest(
                    pdf=test_data["pdf"],
                    page=test_data["page"],
                    id=test_data["id"],
                    type=test_data["type"],
                    cell=test_data["cell"],
                    max_diffs=test_data["max_diffs"],
                    up=test_data.get("up", ""),
                    down=test_data.get("down", ""),
                    left=test_data.get("left", ""),
                    right=test_data.get("right", ""),
                    top_heading=test_data.get("top_heading", ""),
                    left_heading=test_data.get("left_heading", ""),
                )

                # Extract just the relevant table HTML
                tables = soup.find_all("table")
                if table_idx < len(tables):
                    table_html = str(tables[table_idx])

                    # Run the test against the original HTML
                    passed, explanation = test_obj.run(table_html)
                else:
                    # Shouldn't happen, but handle it gracefully
                    passed = False

                # Only add tests that pass
                if passed:
                    table_tests.append(test_data)

                if len(table_tests) > 25:
                    break

        # Done with inner for loop iterating over cells
        # So add in the bulk of the test cases back in now
        tests.extend(table_tests)

    # Step 3: Generate TextPresenceTests and OrderingTests from markdown content
    # Convert HTML to markdown to get cleaner text for presence and ordering tests
    full_markdown_content = html_to_markdown_with_frontmatter(html_content)

    # Extract language from HTML metadata for wordfreq
    metadata = extract_html_metadata(html_content)
    primary_language = metadata.get("primary_language", "en")

    # Remove any HTML tables from the markdown content
    # Tables can persist in markdown as raw HTML and we want to exclude them
    stripped_markdown_content = re.sub(r"<table[^>]*>.*?</table>", "", full_markdown_content, flags=re.DOTALL | re.IGNORECASE)

    # Remove math equations from this markdown content
    # Remove $$...$$ blocks (display math)
    stripped_markdown_content = re.sub(r"\$\$.*?\$\$", "", stripped_markdown_content, flags=re.DOTALL)
    # Remove \[...\] blocks (display math)
    stripped_markdown_content = re.sub(r"\\\[.*?\\\]", "", stripped_markdown_content, flags=re.DOTALL)
    # Remove \(...\) blocks (inline math)
    stripped_markdown_content = re.sub(r"\\\(.*?\\\)", "", stripped_markdown_content, flags=re.DOTALL)

    # Extract just the content part (after frontmatter)
    markdown_lines = stripped_markdown_content.split("\n")
    content_start_idx = 0

    # Skip frontmatter if present
    if markdown_lines[0] == "---":
        for idx, line in enumerate(markdown_lines[1:], 1):
            if line == "---":
                content_start_idx = idx + 1
                break

    # Get markdown content without frontmatter
    markdown_text = "\n".join(markdown_lines[content_start_idx:]).strip()

    # Parse sentences from markdown content
    sentences = []
    if markdown_text:
        for paragraph in process(markdown_text):
            for sentence in paragraph:
                # Convert token sequence to string and clean it
                sentence_str = ""
                for token in sentence:
                    sentence_str += token.spacing + token.value

                sentence_str = sentence_str.strip()

                if sentence_str:
                    # Skip HTML content that might still be in markdown
                    if not sentence_str.startswith("<") and not sentence_str.endswith(">"):
                        # Skip image placeholders - match any markdown image syntax ![...](...)
                        if re.search(r"!\[.*?\]\(.*?\)", sentence_str):
                            continue

                        # Remove leading # marks (markdown headers)
                        while sentence_str.startswith("#"):
                            sentence_str = sentence_str[1:]
                        sentence_str = sentence_str.strip()

                        # Remove leading "- " for unordered lists
                        if sentence_str.startswith("- "):
                            sentence_str = sentence_str[2:]

                        sentence_str = normalize_text(sentence_str.strip())

                        if sentence_str:  # Only add if there's still content after cleaning
                            sentences.append(sentence_str)

    # Add a few random ordering tests
    all_indexes = list(range(len(sentences)))
    random_gen.shuffle(all_indexes)
    random_pairs = [(all_indexes[i * 2], all_indexes[i * 2 + 1]) for i in range(len(all_indexes) // 2)]
    random_pairs = [(min(i, j), max(i, j)) for (i, j) in random_pairs]

    num_order_tests = 0
    for i, j in random_pairs:
        first_sentence = sentences[i]
        second_sentence = sentences[j]

        if len(first_sentence) < 5 or len(second_sentence) < 5:
            continue

        if "\n" in first_sentence:
            first_sentence = first_sentence.split("\n")[0].strip()
        if "\n" in second_sentence:
            second_sentence = second_sentence.split("\n")[0].strip()

        max_diffs = round(max(len(first_sentence), len(second_sentence)) * 0.02)

        # Too big of a length discrepancy causes issues
        if max_diffs > len(first_sentence) // 4 or max_diffs > len(second_sentence) // 4:
            continue

        test_data = {
            "pdf": pdf_filename,
            "page": 1,
            "id": f"{pdf_id}_order_{uuid.uuid4().hex[:8]}",
            "type": TestType.ORDER.value,
            "before": first_sentence,
            "after": second_sentence,
            "max_diffs": max_diffs,
        }

        # Create test object to validate
        try:
            test_obj = TextOrderTest(**test_data)

            # Run the test against the markdown content
            passed, _ = test_obj.run(full_markdown_content)

            if passed:
                num_order_tests += 1
                tests.append(test_data)
        except Exception:
            # Skip if test creation or validation fails
            pass

        if num_order_tests > 5:
            break

    # Go through the top 10 rarest words (any chars permitted), and also the top 5 rarest numbers (only numeric chars allowed)
    # And add those as TestType.PRESENT.value, and also run the test to check that it passes before adding

    word_counter = Counter()
    number_counter = Counter()
    word_rarities = {}  # Store Zipf frequencies for words

    # Split on whitespace and check each token
    # We use the clean set of sentences used in the order test creation, which have been line-by-line normalized
    tokens = "\n".join(sentences).split()

    # Pattern for numbers: optional minus, digits, optional decimal point and more digits
    number_pattern = re.compile(r"^-?\d+(?:\.\d+)?$")

    for token in tokens:
        # Strip common punctuation from ends
        token_cleaned = token.strip(".,;:!?\"'()")

        if not token_cleaned:
            continue

        # Check if it's a number
        if number_pattern.match(token_cleaned):
            if len(token_cleaned) >= 2:  # Only count numbers with at least 2 characters
                number_counter[token_cleaned] += 1
        # Check if it's a word (has at least one alphabetic character)
        elif any(c.isalpha() for c in token_cleaned):
            token_lower = token_cleaned.lower()
            if len(token_lower) >= 4:  # Only count words with at least 4 characters
                word_counter[token_lower] += 1

                # Calculate Zipf frequency if wordfreq is available
                if token_lower not in word_rarities:
                    # Get Zipf frequency (0-8 scale, higher = more common)
                    # Use primary_language if it's a valid 2-letter code
                    lang_code = primary_language if len(primary_language) == 2 else "en"
                    try:
                        zipf = zipf_frequency(token_lower, lang_code)
                        word_rarities[token_lower] = zipf
                    except:
                        # If language not supported or error, use default
                        word_rarities[token_lower] = 4.0  # Assume moderately common

    # Collect rarest items first
    rarest_items = []

    if word_rarities:
        # Filter for truly rare words (Zipf frequency ≤ 3)
        rare_words_with_zipf = [(word, zipf) for word, zipf in word_rarities.items() if zipf <= 3.0]

        # Sort by Zipf frequency (ascending = rarest first), then by count
        rare_words_with_zipf.sort(key=lambda x: (x[1], word_counter[x[0]]))

        # Take up to 10 rarest words
        rarest_words = [word for word, _ in rare_words_with_zipf[:10]]
    else:
        # Fallback to old method if wordfreq not available
        # Get the 10 least common words
        sorted_words = sorted(word_counter.items(), key=lambda x: x[1])
        rarest_words = [word for word, _ in sorted_words[:10]]

    # Shuffle the rarest words
    random_gen.shuffle(rarest_words)

    # Add rarest words to the list with their type
    for word in rarest_words:
        rarest_items.append((word, "present_word"))

    # Get the 5 least common numbers
    sorted_numbers = sorted(number_counter.items(), key=lambda x: x[1])
    rarest_numbers = [num for num, _ in sorted_numbers[:5]]

    # Shuffle the rarest numbers
    random_gen.shuffle(rarest_numbers)

    # Add rarest numbers to the list with their type
    for number in rarest_numbers:
        rarest_items.append((number, "present_num"))

    # Now create and validate presence tests for all rare items
    for test_text, test_id_prefix in rarest_items:
        # Normalize the test text
        normalized_text = normalize_text(test_text)

        if not normalized_text or len(normalized_text) < 2:
            continue

        # Create test data
        test_data = {
            "pdf": pdf_filename,
            "page": 1,
            "id": f"{pdf_id}_{test_id_prefix}_{uuid.uuid4().hex[:8]}",
            "type": TestType.PRESENT.value,
            "text": test_text,
            "max_diffs": 0,
        }

        # Create test object to validate
        try:
            test_obj = TextPresenceTest(
                pdf=test_data["pdf"],
                page=test_data["page"],
                id=test_data["id"],
                type=test_data["type"],
                text=test_data["text"],
                max_diffs=test_data["max_diffs"],
            )

            # Run the test against the markdown content
            passed, _ = test_obj.run(full_markdown_content)

            if passed:
                tests.append(test_data)
        except Exception:
            # Skip if test creation or validation fails
            pass

    # Step 3.5: Generate absence tests for 3 randomly selected common words that don't appear on the page
    # Get the top 1000 most common words using wordfreq
    # Note from Jake: For now, I am commenting this out, I can't just help shake this feeling that this would
    # incentivize the model in slightly the wrong way somehow.
    # lang_code = primary_language if len(primary_language) == 2 else 'en'
    # try:
    #     # Get top 1000 common words
    #     common_words_list = top_n_list(lang_code, 1000)
    # except:
    #     # Fallback to English if language not supported
    #     common_words_list = top_n_list('en', 1000)

    # # Build a set of words that appear on the page (lowercase)
    # page_words_set = set(word.lower() for word in word_counter.keys())

    # # Find common words not on the page
    # absent_common_words = []
    # for word in common_words_list:
    #     if word.lower() not in page_words_set and len(word) >= 2:  # Skip if word is on page
    #         # Get the Zipf frequency to ensure it's truly common
    #         try:
    #             zipf = zipf_frequency(word, lang_code)
    #             absent_common_words.append((word, zipf))
    #         except:
    #             # If error getting Zipf frequency, still include with a high assumed frequency
    #             absent_common_words.append((word, 6.0))

    # # Select 3 absent common words that are most similar to words on the page
    # if len(absent_common_words) > 0 and len(page_words_set) > 0:
    #     from rapidfuzz import fuzz

    #     # Calculate similarity scores for each absent word
    #     absent_words_with_similarity = []
    #     page_words_list = list(page_words_set)  # Convert to list for iteration

    #     for word, zipf_freq in absent_common_words:
    #         # Calculate max similarity to any word on the page
    #         max_similarity = 0
    #         for page_word in page_words_list:
    #             similarity = fuzz.ratio(word.lower(), page_word)
    #             if similarity > max_similarity:
    #                 max_similarity = similarity

    #         absent_words_with_similarity.append((word, zipf_freq, max_similarity))

    #     # Sort by similarity score (descending) to get words most similar to page content
    #     absent_words_with_similarity.sort(key=lambda x: x[2], reverse=True)

    #     # Select top 3 most similar words
    #     num_to_select = min(3, len(absent_words_with_similarity))
    #     selected_absent_words = absent_words_with_similarity[:num_to_select]

    #     # Create absence tests for these selected common words
    #     for word, zipf_freq, similarity_score in selected_absent_words:
    #         test_data = {
    #             "pdf": pdf_filename,
    #             "page": 1,
    #             "id": f"{pdf_id}_absent_common_{uuid.uuid4().hex[:8]}",
    #             "type": TestType.ABSENT.value,
    #             "text": word,
    #             "max_diffs": 0,
    #             "case_sensitive": False,
    #         }

    #         # Double-check the word really doesn't appear in the markdown text
    #         # For ABSENT tests, we want to ensure the word is NOT found
    #         try:
    #             validation_test = TextPresenceTest(**test_data)

    #             # Run the validation test against the markdown content
    #             passed, _ = validation_test.run(markdown_text)

    #             if passed:
    #                 tests.append(test_data)
    #         except Exception:
    #             # Skip if test creation or validation fails
    #             pass

    # Step 4: Generate Math tests for LaTeX equations from the markdown

    # Define math patterns to search for
    math_patterns = [
        (r"\$\$(.+?)\$\$", re.DOTALL),  # $$...$$ (multiline)
        (r"\\\((.+?)\\\)", re.DOTALL),  # \(...\) (multiline)
        (r"\\\[(.+?)\\\]", re.DOTALL),  # \[...\] (multiline)
    ]

    math_equations = []
    for pattern, flags in math_patterns:
        matches = re.findall(pattern, full_markdown_content, flags)
        for match in matches:
            # Clean up the match - remove extra whitespace and newlines
            equation = match.strip()
            # Skip empty or very short equations
            if len(equation) > 2:
                math_equations.append(equation)

    # Remove duplicates while preserving order
    seen = set()
    unique_equations = []
    for eq in math_equations:
        if eq not in seen:
            seen.add(eq)
            unique_equations.append(eq)

    # Create math tests for up to 50 unique equations
    for i, equation in enumerate(unique_equations[:50]):
        tests.append(
            {
                "pdf": pdf_filename,
                "page": 1,
                "id": f"{pdf_id}_math_{uuid.uuid4().hex[:8]}",
                "type": "math",
                "math": equation,
                "max_diffs": 0,
                "ignore_dollar_delimited": True,
            }
        )

    # Step 5: Generate FormatTests for headings, bold, and italic text
    format_tests = []

    # Define mapping from HTML tags to format types
    format_tag_mapping = {("h1", "h2", "h3", "h4", "h5", "h6"): "heading", ("b", "strong"): "bold", ("i", "em"): "italic"}

    # Parse the HTML to find formatted elements
    format_soup = BeautifulSoup(html_content, "html.parser")

    # Convert superscripts and subscripts before extracting text
    convert_superscripts_subscripts(format_soup)

    # Remove headers, footers, and tables from format soup to focus on main content
    for element in format_soup.find_all(["header", "footer", "table"]):
        element.decompose()

    # Track what we've already tested to avoid duplicates
    tested_texts = set()

    # Process each format type
    for tags, format_type in format_tag_mapping.items():
        # Find all elements with these tags
        elements = format_soup.find_all(list(tags))

        for element in elements:
            if len(format_tests) >= 5:  # Limit to 5 total format tests
                break

            element_text = element.get_text().strip()
            element_text = normalize_text(element_text)

            if element_text and len(element_text) >= 3 and element_text not in tested_texts:
                # Create a format test
                test_data = {
                    "pdf": pdf_filename,
                    "page": 1,
                    "id": f"{pdf_id}_format_{format_type}_{uuid.uuid4().hex[:8]}",
                    "type": TestType.FORMAT.value,
                    "text": element_text,
                    "format": format_type,
                    "max_diffs": round(len(element_text) * 0.05),
                }

                # Validate the test against markdown_content
                try:
                    test_obj = FormatTest(
                        pdf=test_data["pdf"],
                        page=test_data["page"],
                        id=test_data["id"],
                        type=test_data["type"],
                        text=test_data["text"],
                        format=test_data["format"],
                        max_diffs=test_data["max_diffs"],
                    )

                    # Test against the markdown_content
                    passed, _ = test_obj.run(full_markdown_content)

                    if passed:
                        format_tests.append(test_data)
                        tested_texts.add(element_text)
                except Exception:
                    pass

    # Add format tests to the main tests list
    tests.extend(format_tests)

    # Step 6: Generate FootnoteTests for footnotes on the page
    footnote_tests = []

    # Parse the HTML to find footnotes (without converting superscripts)
    footnote_soup = BeautifulSoup(html_content, "html.parser")
    cleanup_headers_footers_soup(footnote_soup)

    # Look for superscript elements that might be footnote markers
    sup_elements = footnote_soup.find_all("sup")

    max_footnote_tests = 5
    marker_sup_map = {}

    for sup in sup_elements:
        marker_text = sup.get_text().strip()

        # Filter out markers that are unlikely to be footnote references
        if not marker_text or not (marker_text.isdigit() or (len(marker_text) == 1 and marker_text.isalpha()) or marker_text in ["*", "†", "‡", "§", "¶"]):
            continue

        if marker_text not in marker_sup_map:
            if len(marker_sup_map) >= max_footnote_tests:
                continue
            marker_sup_map[marker_text] = []

        marker_sup_map[marker_text].append(sup)

    for marker_text, marker_superscripts in marker_sup_map.items():
        test_data = {
            "pdf": pdf_filename,
            "page": 1,
            "id": f"{pdf_id}_footnote_{uuid.uuid4().hex[:8]}",
            "type": TestType.FOOTNOTE.value,
            "marker": marker_text,
            "max_diffs": 0,
        }

        # Extract text context for each occurrence
        occurrences = []
        for sup in marker_superscripts:
            parent = sup.parent
            if not parent:
                continue

            # Get text before and after the sup element by reconstructing parent content
            before_text = None
            after_text = None

            # Get all content before this sup element
            text_before_sup = []
            for sibling in parent.children:
                if sibling == sup:
                    break
                if isinstance(sibling, str):
                    text_before_sup.append(sibling)
                else:
                    text_before_sup.append(sibling.get_text())

            # Get all content after this sup element
            found_sup = False
            text_after_sup = []
            for sibling in parent.children:
                if found_sup:
                    if isinstance(sibling, str):
                        text_after_sup.append(sibling)
                    else:
                        text_after_sup.append(sibling.get_text())
                elif sibling == sup:
                    found_sup = True

            # Process text before marker
            preceding_text = "".join(text_before_sup).strip()
            if len(preceding_text) >= 10:
                words = preceding_text.split()
                if len(words) >= 2:
                    last_words = " ".join(words[-3:]) if len(words) >= 3 else " ".join(words)
                else:
                    last_words = preceding_text
                candidate = normalize_text(last_words)
                if candidate and len(candidate) >= 5:
                    before_text = candidate

            # Process text after marker
            following_text = "".join(text_after_sup).strip()
            if len(following_text) >= 10:
                words = following_text.split()
                if len(words) >= 2:
                    first_words = " ".join(words[:3]) if len(words) >= 3 else " ".join(words)
                else:
                    first_words = following_text[:50]
                candidate = normalize_text(first_words)
                if candidate and len(candidate) >= 5:
                    after_text = candidate

            occurrences.append({"before": before_text, "after": after_text})

        # Apply logic based on number of occurrences
        if len(occurrences) >= 2:
            # If marker exists 2+ times: use before from first, after from second
            if occurrences[0]["before"]:
                test_data["appears_before_marker"] = occurrences[0]["before"]
            if occurrences[1]["after"]:
                test_data["appears_after_marker"] = occurrences[1]["after"]
        elif len(occurrences) == 1:
            # If marker exists 1 time: prefer left text, then right text
            if occurrences[0]["before"]:
                test_data["appears_before_marker"] = occurrences[0]["before"]
            elif occurrences[0]["after"]:
                test_data["appears_after_marker"] = occurrences[0]["after"]
            # If neither before nor after text exists, we'll just have the marker

        # Create test even if we only have the marker (no additional fields required)
        try:
            test_obj = FootnoteTest(**test_data)
            passed, _ = test_obj.run(full_markdown_content)
            if passed:
                footnote_tests.append(test_data)
        except Exception:
            pass

    # Add footnote tests to the main tests list
    tests.extend(footnote_tests)

    # Final test filtering out stage

    # Now double check that the absent tests don't find any matches in the markdown_text
    # If they do, filter them out
    tests = [t for t in tests if t["type"] != "absent" or t["text"] not in markdown_text]

    # Remove any tests where text-based fields have no alphanumeric characters, contain LaTeX, or contain Unicode super/subscripts
    text_fields = ["text", "cell", "before", "after", "up", "down", "left", "right", "top_heading", "left_heading"]

    def contains_alphanumeric(value):
        return any(c.isalnum() for c in value) if isinstance(value, str) else False

    def contains_latex(value):
        if not isinstance(value, str):
            return False
        # Check for LaTeX delimiters
        latex_patterns = [r"\(", r"\)", r"\[", r"\]"]
        return any(pattern in value for pattern in latex_patterns)

    def contains_unicode_super_or_subscripts(value):
        if not isinstance(value, str):
            return False

        # Unicode ranges for superscripts and subscripts
        superscript_chars = "⁰¹²³⁴⁵⁶⁷⁸⁹⁺⁻⁼⁽⁾ⁿⁱ"
        subscript_chars = "₀₁₂₃₄₅₆₇₈₉₊₋₌₍₎ₐₑₒₓₕₖₗₘₙₚₛₜ"

        return any(c in superscript_chars or c in subscript_chars for c in value)

    def contains_html_sup_sub_tags(value):
        if not isinstance(value, str):
            return False
        # Check for HTML sup/sub tags
        return "<sup>" in value or "</sup>" in value or "<sub>" in value or "</sub>" in value

    filtered_tests = []
    for test in tests:
        # Math tests should not be filtered for LaTeX content
        if test.get("type") == "math":
            filtered_tests.append(test)
            continue

        # Format tests have different validation requirements
        if test.get("type") == TestType.FORMAT.value:
            # Only check the "text" field for format tests
            if "text" in test:
                if contains_alphanumeric(test["text"]) and not contains_latex(test["text"]) and not contains_unicode_super_or_subscripts(test["text"]):
                    filtered_tests.append(test)
            continue

        # Footnote tests have special requirements
        if test.get("type") == TestType.FOOTNOTE.value:
            # Markers can contain superscript characters, so don't filter them
            # But appears_before_marker and appears_after_marker should not contain LaTeX
            valid = True
            if "appears_before_marker" in test and test["appears_before_marker"]:
                if not contains_alphanumeric(test["appears_before_marker"]) or contains_latex(test["appears_before_marker"]):
                    valid = False
            if "appears_after_marker" in test and test["appears_after_marker"]:
                if not contains_alphanumeric(test["appears_after_marker"]) or contains_latex(test["appears_after_marker"]):
                    valid = False
            if valid:
                filtered_tests.append(test)
            continue

        # Check all text fields in the test for alphanumeric content, LaTeX, Unicode super/subscripts, and HTML tags
        all_valid = True
        for field in text_fields:
            if field in test:
                # Skip test if field has no alphanumeric characters
                if not contains_alphanumeric(test[field]):
                    all_valid = False
                    break
                # Skip test if field contains LaTeX delimiters
                if contains_latex(test[field]):
                    all_valid = False
                    break
                # Skip test if field contains Unicode super or subscripts
                if contains_unicode_super_or_subscripts(test[field]):
                    all_valid = False
                    break
                # Skip test if field contains HTML sup/sub tags
                if contains_html_sup_sub_tags(test[field]):
                    all_valid = False
                    break
        if all_valid:
            filtered_tests.append(test)

    tests = filtered_tests

    # Remove duplicate tests (identical on everything but the id field)
    unique_tests = []
    test_signatures = set()

    for test in tests:
        # Create a signature for the test by using all fields except 'id'
        test_dict = test.copy()
        test_dict.pop("id")

        # Convert dict to a sorted tuple of items for hashability
        test_signature = tuple(sorted((k, str(v)) for k, v in test_dict.items()))

        # Only add the test if we haven't seen an identical one
        if test_signature not in test_signatures:
            test_signatures.add(test_signature)
            unique_tests.append(test)

    # Add a single BaselineTest for this page
    # First, check if the markdown content would fail due to disallowed characters
    baseline_test_data = {
        "pdf": pdf_filename,
        "page": 1,
        "id": f"{pdf_id}_baseline_{uuid.uuid4().hex[:8]}",
        "type": "baseline",
        "max_repeats": 30,
        "check_disallowed_characters": True,
    }

    # Create a baseline test object to check if disallowed characters are present
    try:
        baseline_test_obj = BaselineTest(
            pdf=baseline_test_data["pdf"],
            page=baseline_test_data["page"],
            id=baseline_test_data["id"],
            type=baseline_test_data["type"],
            max_repeats=baseline_test_data["max_repeats"],
            check_disallowed_characters=True,
        )

        # Run the test with check_disallowed_characters=True
        passed, explanation = baseline_test_obj.run(full_markdown_content)

        # If it failed due to disallowed characters, set check_disallowed_characters=False
        if not passed and "disallowed characters" in explanation:
            baseline_test_data["check_disallowed_characters"] = False

    except Exception:
        # If there was an error creating/running the test, keep default settings
        pass

    # Add the baseline test to the list
    unique_tests.append(baseline_test_data)

    return unique_tests


def check_outputs_exist(pdf_id: str, page_num: int, args, existing_test_pdfs: set) -> bool:
    """
    Check if all output files for a given PDF already exist.

    Returns True if all outputs exist and processing can be skipped.
    """
    # Define expected output paths
    html_dir = os.path.join(args.output_dir, "html", args.name)
    pdfs_dir = os.path.join(args.output_dir, "pdfs", args.name)
    training_dir = os.path.join(args.output_dir, "training", args.name)
    bench_data_dir = os.path.join(args.output_dir, "bench_data")
    claude_original_dir = os.path.join(bench_data_dir, "claude_original", args.name)

    base_filename = f"{pdf_id}_page{page_num}"

    # Check HTML file
    html_path = os.path.join(html_dir, f"{base_filename}.html")
    if not os.path.exists(html_path):
        return False

    # Check rendered PDF in pdfs dir
    pdf_path = os.path.join(pdfs_dir, f"{base_filename}.pdf")
    if not os.path.exists(pdf_path):
        return False

    # Check markdown in training dir
    markdown_path = os.path.join(training_dir, f"{base_filename}.md")
    if not os.path.exists(markdown_path):
        return False

    # Check symlink in training dir
    pdf_link_path = os.path.join(training_dir, f"{base_filename}.pdf")
    if not os.path.exists(pdf_link_path) and not os.path.islink(pdf_link_path):
        return False

    # Check claude_original symlink
    claude_md_link_path = os.path.join(claude_original_dir, f"{base_filename}_pg1_repeat1.md")
    if not os.path.exists(claude_md_link_path) and not os.path.islink(claude_md_link_path):
        return False

    # Check that bench data has at least one test for this PDF
    expected_test_pdf = f"{args.name}/{base_filename}.pdf"
    if expected_test_pdf not in existing_test_pdfs:
        return False

    return True


def load_existing_test_pdfs(args) -> set:
    """
    Load existing test PDFs from the bench data JSONL file.

    Returns a set of PDF paths that already have tests.
    """
    bench_data_dir = os.path.join(args.output_dir, "bench_data")
    synthetic_json_path = os.path.join(bench_data_dir, f"{args.name}.jsonl")

    existing_pdfs = set()
    if os.path.exists(synthetic_json_path):
        try:
            with open(synthetic_json_path, "r") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            test = json.loads(line)
                            if "pdf" in test:
                                existing_pdfs.add(test["pdf"])
                        except json.JSONDecodeError:
                            continue
        except Exception as e:
            print(f"Warning: Could not read existing bench data: {e}")

    return existing_pdfs


async def process_pdf(pdf_info, args, client, pdf_filter=None, existing_test_pdfs=None):
    """Process a single PDF, render a random page, and create an HTML template."""
    pdf_path, index = pdf_info

    # Create a unique folder for each PDF in the temp directory; include PID and PDF path hash to avoid cross-run collisions
    pdf_id = f"pdf_{index:05d}"
    temp_dir_suffix = f"{os.getpid()}_{hashlib.sha1(pdf_path.encode('utf-8')).hexdigest()[:16]}"
    temp_pdf_dir = os.path.join(args.temp_dir, f"{pdf_id}_{temp_dir_suffix}")
    os.makedirs(temp_pdf_dir, exist_ok=True)

    # Determine if we should log table test verification
    verbose_table_testing = args.verbose

    # Download PDF to local temp directory (or copy if local)
    local_pdf_path = os.path.join(temp_pdf_dir, "document.pdf")
    if not download_s3_pdf(pdf_path, local_pdf_path):
        print(f"Failed to download/copy PDF from {pdf_path}")
        return None

    # Apply filter if enabled
    if pdf_filter and pdf_filter.filter_out_pdf(local_pdf_path):
        print(f"PDF filtered out: {pdf_path}")
        return None

    # Seed with SHA1 hash of PDF contents for reproducibility
    with open(local_pdf_path, "rb") as f:
        pdf_content = f.read()
        pdf_hash = hashlib.sha1(pdf_content).hexdigest()

    # Use the first 8 characters of the hash as an integer seed
    seed = int(pdf_hash[:8], 16)
    random_generator = random.Random(seed)

    try:
        # Get page count using pypdf
        reader = pypdf.PdfReader(local_pdf_path)
        num_pages = len(reader.pages)

        if num_pages == 0:
            print(f"PDF has no pages: {pdf_path}")
            return None

        # Select a random page
        page_num = random_generator.randint(1, num_pages)

        # Check if this PDF has already been processed (skip logic for resume)
        if existing_test_pdfs is not None and check_outputs_exist(pdf_id, page_num, args, existing_test_pdfs):
            print(f"Skipping {pdf_id} page {page_num} - already processed")
            return {"skipped": True, "pdf_id": pdf_id, "page_number": page_num}

        # Render the page as a base64 PNG (run in thread pool since it's blocking I/O)
        loop = asyncio.get_event_loop()
        image_base64 = await loop.run_in_executor(None, render_pdf_to_base64png, local_pdf_path, page_num, 1024)

        # Generate HTML from the image
        html_content = await generate_html_from_image(client, image_base64)
        if not html_content:
            print(f"Failed to generate HTML for {pdf_path}, page {page_num}")
            return None

        if args.densify:
            from olmocr.synth.augmentations import densify_html

            html_content = await densify_html(client, html_content)

            if not html_content:
                print(f"Failed to densify HTML for {pdf_path}, page {page_num}")
                return None

        typo_records = []
        pre_typo_html = None
        if args.introduce_text_errors > 0:
            from olmocr.synth.augmentations import introduce_text_errors

            pre_typo_html = html_content
            html_content, typo_records = introduce_text_errors(html_content, random_generator, num_errors=args.introduce_text_errors)

        # Add git commit meta tag if available
        git_commit = get_git_commit_hash()
        if git_commit:
            # Parse the HTML to add the meta tag in the head section
            html_soup = BeautifulSoup(html_content, "html.parser")

            # Only add meta tag if head element exists
            head = html_soup.find("head")
            if head:
                # Add meta tag with git commit
                meta_tag = html_soup.new_tag("meta", attrs={"name": "olmocr_git_commit", "content": git_commit})
                head.append(meta_tag)

                # Update initial_html with the modified version
                html_content = str(html_soup)

        # Create output directories
        html_dir = os.path.join(args.output_dir, "html", args.name)
        pdfs_dir = os.path.join(args.output_dir, "pdfs", args.name)
        training_dir = os.path.join(args.output_dir, "training", args.name)
        bench_data_dir = os.path.join(args.output_dir, "bench_data")
        bench_synthetic_dir = os.path.join(bench_data_dir, "pdfs", args.name)
        claude_original_dir = os.path.join(bench_data_dir, "claude_original", args.name)
        os.makedirs(html_dir, exist_ok=True)
        os.makedirs(pdfs_dir, exist_ok=True)
        os.makedirs(training_dir, exist_ok=True)
        os.makedirs(bench_data_dir, exist_ok=True)
        os.makedirs(bench_synthetic_dir, exist_ok=True)
        os.makedirs(claude_original_dir, exist_ok=True)

        # Render PDF using Playwright
        playwright_pdf_path = None
        render_result = None
        playwright_pdf_filename = f"{pdf_id}_page{page_num}.pdf"  # This will be used in the tests

        playwright_pdf_path = os.path.join(pdfs_dir, playwright_pdf_filename)

        try:
            # Get PNG dimensions
            png_width, png_height = get_png_dimensions_from_base64(image_base64)

            # Run the async function directly since we're already in an async context
            render_result = await render_pdf_with_playwright(html_content, playwright_pdf_path, png_width, png_height)

            if render_result.success:
                print(f"Successfully rendered with Playwright: {playwright_pdf_path}")

                # Apply JPEG compression if requested
                if args.jpegify:
                    # Select a random quality level between 70 and 95
                    jpeg_quality = random_generator.randint(70, 95)
                    print(f"Applying JPEG compression with quality {jpeg_quality} to {playwright_pdf_path}")

                    from olmocr.synth.augmentations import apply_jpeg_compression

                    compression_success = apply_jpeg_compression(playwright_pdf_path, jpeg_quality, temp_pdf_dir)

                    if compression_success:
                        print(f"Successfully applied JPEG compression with quality {jpeg_quality}")
                    else:
                        print(f"Warning: Failed to apply JPEG compression, keeping original PDF")

            else:
                if render_result.has_cutoff:
                    print(f"Skipping: text cutoff detected in {playwright_pdf_path}")
                else:
                    print(f"Failed to render as a single page PDF: {playwright_pdf_path}")
                playwright_pdf_path = None
        except Exception as e:
            print(f"Failed to render with Playwright: {e}")
            playwright_pdf_path = None
            render_result = None

        # If playwright rendering failed and was required, return None to skip the rest of the output here
        if not render_result or not render_result.success:
            return None

        # Save HTML to output directory
        html_path = os.path.join(html_dir, f"{pdf_id}_page{page_num}.html")
        with open(html_path, "w") as f:
            f.write(html_content)

        # Convert HTML to markdown with FrontMatter and save
        markdown_content = html_to_markdown_with_frontmatter(html_content)
        markdown_filename = f"{pdf_id}_page{page_num}.md"
        markdown_path = os.path.join(training_dir, markdown_filename)
        with open(markdown_path, "w") as f:
            f.write(markdown_content)

        # Create soft link to PDF in training directory
        pdf_link_name = f"{pdf_id}_page{page_num}.pdf"
        pdf_link_path = os.path.join(training_dir, pdf_link_name)
        # Remove existing link if it exists
        if os.path.exists(pdf_link_path) or os.path.islink(pdf_link_path):
            os.remove(pdf_link_path)
        # Create relative symlink from training to pdfs directory
        os.symlink(os.path.relpath(os.path.join(pdfs_dir, f"{pdf_id}_page{page_num}.pdf"), training_dir), pdf_link_path)

        # Create soft link to markdown in claude_original/synthetic with new naming scheme
        claude_md_link_name = f"{pdf_id}_page{page_num}_pg1_repeat1.md"
        claude_md_link_path = os.path.join(claude_original_dir, claude_md_link_name)
        # Remove existing link if it exists
        if os.path.exists(claude_md_link_path) or os.path.islink(claude_md_link_path):
            os.remove(claude_md_link_path)
        # Create relative symlink from claude_original/synthetic to training directory
        os.symlink(os.path.relpath(markdown_path, claude_original_dir), claude_md_link_path)

        # Extract the page and save as PDF
        original_pdf_path = os.path.join(pdfs_dir, f"{pdf_id}_page{page_num}_original.pdf")
        if not extract_page_from_pdf(local_pdf_path, original_pdf_path, page_num):
            print(f"Failed to extract page {page_num} from {local_pdf_path}")

        # Create soft link in bench_data/synthetic/ directory
        if playwright_pdf_path:
            synthetic_link_path = os.path.join(bench_synthetic_dir, playwright_pdf_filename)
            # Remove existing link if it exists
            if os.path.exists(synthetic_link_path) or os.path.islink(synthetic_link_path):
                os.remove(synthetic_link_path)
            # Create relative symlink from bench_data/synthetic to pdfs directory
            os.symlink(os.path.relpath(playwright_pdf_path, bench_synthetic_dir), synthetic_link_path)

        # Generate tests from the HTML content
        # Use the playwright rendered PDF path for tests
        tests = generate_tests_from_html(html_content, pdf_id, page_num, random_generator, verbose_table_testing)

        # Update the PDF path in all tests to use the playwright rendered PDF with the specified name prefix
        for test in tests:
            test["pdf"] = f"{args.name}/{playwright_pdf_filename}"

        if args.introduce_text_errors > 0 and typo_records and pre_typo_html is not None:
            original_markdown = html_to_markdown_with_frontmatter(pre_typo_html)
            augmented_markdown = html_to_markdown_with_frontmatter(html_content)

            # Remove any existing presence tests that match typo words
            # (generate_tests_from_html may tag them as rare words)
            typo_words = {r["typo_word"] for r in typo_records}
            tests = [t for t in tests if not (t.get("type") == TestType.PRESENT.value and t.get("text") in typo_words)]

            for record in typo_records:
                typo_word = record["typo_word"]

                test_id = f"{pdf_id}_typo_{uuid.uuid4().hex[:8]}"
                test_data = {
                    "pdf": f"{args.name}/{playwright_pdf_filename}",
                    "page": 1,
                    "id": test_id,
                    "type": TestType.PRESENT.value,
                    "text": typo_word,
                    "max_diffs": 0,
                }

                test_obj = TextPresenceTest(
                    pdf=test_data["pdf"],
                    page=1,
                    id=test_id,
                    type=TestType.PRESENT.value,
                    text=typo_word,
                    max_diffs=0,
                )

                # Must FAIL against original (typo shouldn't exist before)
                passed_original, _ = test_obj.run(original_markdown)
                # Must PASS against augmented (typo should exist after)
                passed_augmented, _ = test_obj.run(augmented_markdown)

                if not passed_original and passed_augmented:
                    tests.append(test_data)

        # Log table test stats if verbose
        if verbose_table_testing:
            table_tests = [t for t in tests if t["type"] == TestType.TABLE.value]
            print(f"Generated {len(table_tests)} table tests for {pdf_id}, page {page_num} (passed verification)")

        return {
            "pdf_id": pdf_id,
            "pdf_path": pdf_path,
            "page_number": page_num,
            "html_path": html_path,
            "markdown_path": markdown_path,
            "original_pdf_path": original_pdf_path,
            "playwright_pdf_path": playwright_pdf_path,
            "tests": tests,
            "num_tests": len(tests),
        }
    except Exception as e:
        print(f"Error processing {pdf_path}: {e}")
        return None
    finally:
        # Clean up temp directory for this PDF
        if os.path.exists(temp_pdf_dir):
            subprocess.run(["rm", "-rf", temp_pdf_dir])


async def main():
    # Configure logging to suppress httpx messages
    logging.getLogger("httpx").setLevel(logging.WARNING)
    logging.getLogger("httpcore").setLevel(logging.WARNING)

    parser = argparse.ArgumentParser(description="Convert PDFs to HTML templates and render with Playwright")
    parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths or local paths to PDFs")
    parser.add_argument("--output_dir", required=True, help="Directory to store extracted pages and tests")
    parser.add_argument("--temp_dir", default="/tmp/mine_tables", help="Directory for temporary files")
    parser.add_argument("--max_tests", type=int, default=100, help="Maximum number of tests to generate")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for sampling selection of PDFs")
    parser.add_argument("--parallel", type=int, default=1, help="Number of parallel tasks to use")
    parser.add_argument("--api_key", help="Claude API key (or set ANTHROPIC_API_KEY environment variable)")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose output including table test verification")
    parser.add_argument("--densify", action="store_true", help="Set to ask claude to double the density of information on this page synthetically")
    parser.add_argument("--jpegify", action="store_true", help="Apply JPEG compression to rendered PDFs with random quality (70-95)")
    parser.add_argument(
        "--introduce-text-errors", type=int, default=0, help="Introduce N intentional typos into HTML body text and generate corresponding presence tests"
    )
    parser.add_argument("--filter", action="store_true", help="Apply PDF filtering to remove forms, spam, and non-English content")
    parser.add_argument("--name", default="synthetic", help="Name for the output JSONL file and subfolder (default: synthetic)")
    args = parser.parse_args()

    # Ensure output and temp directories exist
    os.makedirs(args.output_dir, exist_ok=True)
    os.makedirs(args.temp_dir, exist_ok=True)

    # Get API key
    api_key = args.api_key or os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print("Error: API key not provided. Use --api_key or set ANTHROPIC_API_KEY environment variable.")
        return

    # Initialize async Claude client
    client = AsyncAnthropic(api_key=api_key)

    # Initialize PDF filter if enabled
    pdf_filter = None
    if args.filter:
        pdf_filter = PdfFilter(
            languages_to_keep={Language.ENGLISH, None},  # None means could not detect language, that's okay keep it, might be an OCR
            apply_download_spam_check=True,
            apply_form_check=True,
        )
        print("PDF filtering enabled")

    # Reservoir sampling implementation
    random_gen = random.Random(args.seed)
    pdf_paths = []

    if os.path.isdir(args.input_list):
        pdf_paths = list(sorted(glob.glob(os.path.join(args.input_list, "*.pdf"), recursive=True)))
    else:
        with open(args.input_list, "r") as f:
            for i, line in enumerate(tqdm(f)):
                line = line.strip()
                if not line:
                    continue

                if i < 100000:
                    pdf_paths.append(line)
                else:
                    # Randomly replace elements with decreasing probability
                    j = random_gen.randint(0, i)
                    if j < 100000:
                        pdf_paths[j] = line

    print(f"Found {len(pdf_paths)} PDF paths in input list")

    # Shuffle and limit to max_tests
    random_gen.shuffle(pdf_paths)
    pdf_paths = pdf_paths[: args.max_tests]

    # Initialize the JSONL file in bench_data folder with the specified name
    bench_data_dir = os.path.join(args.output_dir, "bench_data")
    os.makedirs(bench_data_dir, exist_ok=True)
    synthetic_json_path = os.path.join(bench_data_dir, f"{args.name}.jsonl")

    # Load existing test PDFs for skip logic (resume support)
    existing_test_pdfs = load_existing_test_pdfs(args)
    if existing_test_pdfs:
        print(f"Found {len(existing_test_pdfs)} existing PDFs with tests - will skip already processed items")

    # Initialize the metadata JSONL file
    metadata_dir = os.path.join(args.output_dir, "metadata")
    os.makedirs(metadata_dir, exist_ok=True)
    metadata_json_path = os.path.join(metadata_dir, f"{args.name}.jsonl")

    # Counter for test statistics
    test_counter = 0
    test_types = defaultdict(int)  # Automatically handles any test type
    results = []

    # Tracking for success/failure rates
    total_attempted = 0
    successful_templates = 0
    failed_templates = 0
    skipped_templates = 0
    failure_reasons = defaultdict(int)

    # Initialize an asyncio lock for file access
    file_lock = asyncio.Lock()

    # Process PDFs in parallel using asyncio
    async def process_with_progress(pdf_info):
        pdf_path = pdf_info[0]
        nonlocal total_attempted, successful_templates, failed_templates

        async with file_lock:
            total_attempted += 1

        try:
            result = await process_pdf(pdf_info, args, client, pdf_filter, existing_test_pdfs)

            # Handle skipped results (already processed)
            if result and result.get("skipped"):
                nonlocal skipped_templates
                async with file_lock:
                    skipped_templates += 1
                return result

            if result and result.get("tests"):
                # Append tests to synthetic.json as they're created (JSONL format)
                async with file_lock:
                    # Append each test as a separate JSON line
                    with open(synthetic_json_path, "a") as f:
                        for test in result["tests"]:
                            f.write(json.dumps(test) + "\n")

                    # Write metadata mapping (pdf_id to source URL)
                    with open(metadata_json_path, "a") as f:
                        metadata = {"pdf_id": result["pdf_id"], "source_url": result["pdf_path"], "page_number": result["page_number"]}
                        f.write(json.dumps(metadata) + "\n")

                    # Update counters
                    nonlocal test_counter
                    test_counter += len(result["tests"])
                    for test in result["tests"]:
                        test_type = test.get("type", "unknown")
                        test_types[test_type] += 1

                    successful_templates += 1
                    print(f"Added {len(result['tests'])} tests from {result['pdf_id']}, total: {test_counter}")

                return result
            else:
                async with file_lock:
                    failed_templates += 1
                    if result is None:
                        failure_reasons["processing_failed"] += 1
                    elif not result.get("tests"):
                        failure_reasons["no_tests_generated"] += 1
                return None
        except Exception as e:
            print(f"Error processing {pdf_path}: {e}")
            async with file_lock:
                failed_templates += 1
                failure_reasons["exception"] += 1
            return None

    # Create tasks for all PDFs
    tasks = []
    for i, pdf_path in enumerate(pdf_paths):
        tasks.append(process_with_progress((pdf_path, i)))

    # Run tasks with limited concurrency
    semaphore = asyncio.Semaphore(args.parallel)

    async def bounded_task(task_coro):
        async with semaphore:
            return await task_coro

    bounded_tasks = [bounded_task(task) for task in tasks]

    # Process all tasks with progress bar
    pbar = tqdm(asyncio.as_completed(bounded_tasks), total=len(bounded_tasks), desc="Processing PDFs")
    for coro in pbar:
        result = await coro
        if result:
            results.append(result)

        # Update progress bar with cost information
        cost_input = (total_input_tokens / 1_000_000) * 3.0  # $3 per million input tokens
        cost_output = (total_output_tokens / 1_000_000) * 15.0  # $15 per million output tokens
        total_cost = cost_input + cost_output
        pbar.set_postfix({"in_tokens": f"{total_input_tokens:,}", "out_tokens": f"{total_output_tokens:,}", "cost": f"${total_cost:.2f}"})

    print(f"Generated {len(results)} HTML templates")

    # Print summary of Playwright rendering results
    playwright_success = sum(1 for r in results if r and r.get("playwright_pdf_path"))
    print(f"Playwright PDF rendering: {playwright_success}/{len(results)} successful")

    print(f"Saved {test_counter} tests to {synthetic_json_path}")

    # Print summary of generated tests
    print(f"Generated a total of {test_counter} tests across {len(results)} templates")

    # Print test type distribution
    if test_counter > 0:
        print("Test type distribution:")
        for test_type, count in test_types.items():
            print(f"  - {test_type}: {count} tests")

    # Print failure rate summary
    print("\n===== Template Generation Summary =====")
    print(f"Total PDFs attempted: {total_attempted}")
    if total_attempted > 0:
        print(f"Skipped (already processed): {skipped_templates} ({skipped_templates/total_attempted*100:.1f}%)")
        print(f"Successfully generated: {successful_templates} ({successful_templates/total_attempted*100:.1f}%)")
        print(f"Failed: {failed_templates} ({failed_templates/total_attempted*100:.1f}%)")
    else:
        print("No PDFs were processed")

    if failed_templates > 0 and failure_reasons:
        print("\nFailure breakdown:")
        for reason, count in sorted(failure_reasons.items()):
            print(f"  - {reason}: {count}")

    # Print final Claude API cost summary
    print("\nClaude Sonnet API Usage Summary:")
    print(f"  Total input tokens: {total_input_tokens:,}")
    print(f"  Total output tokens: {total_output_tokens:,}")
    cost_input = (total_input_tokens / 1_000_000) * 3.0
    cost_output = (total_output_tokens / 1_000_000) * 15.0
    total_cost = cost_input + cost_output
    print(f"  Input cost: ${cost_input:.2f} ($3/MTok)")
    print(f"  Output cost: ${cost_output:.2f} ($15/MTok)")
    print(f"  Total cost: ${total_cost:.2f}")


if __name__ == "__main__":
    asyncio.run(main())