rerender_pdf.py

import os
import glob
import argparse
from rich.console import Console
from rich.panel import Panel
from synth import compile_latex_to_pdf

# Initialize rich console
console = Console()


def find_latest_run_directory(output_dir="generated"):
    """Find the most recent run directory in the output folder"""
    try:
        run_dirs = [
            d
            for d in os.listdir(output_dir)
            if os.path.isdir(os.path.join(output_dir, d))
        ]
        run_dirs.sort(reverse=True)
        if not run_dirs:
            console.print(
                "[bold red]No run directories found in the output folder.[/bold red]"
            )
            return None
        return os.path.join(output_dir, run_dirs[0])
    except Exception as e:
        console.print(f"[bold red]Error finding latest run directory: {e}[/bold red]")
        return None


def find_latex_file(run_dir):
    """Find the LaTeX file in the final directory of the run"""
    try:
        final_dir = os.path.join(run_dir, "final")
        if not os.path.exists(final_dir):
            console.print(
                f"[bold red]Final directory not found in {run_dir}[/bold red]"
            )
            return None
        tex_files = glob.glob(os.path.join(final_dir, "*.tex"))
        if not tex_files:
            console.print(f"[bold red]No LaTeX files found in {final_dir}[/bold red]")
            return None
        return tex_files[0]
    except Exception as e:
        console.print(f"[bold red]Error finding LaTeX file: {e}[/bold red]")
        return None


def find_latex_file_in_dir(final_dir):
    """Find the LaTeX file directly in a given directory"""
    try:
        if not os.path.exists(final_dir) or not os.path.isdir(final_dir):
            console.print(
                f"[bold red]Provided directory not found: {final_dir}[/bold red]"
            )
            return None
        tex_files = glob.glob(os.path.join(final_dir, "*.tex"))
        if not tex_files:
            console.print(f"[bold red]No LaTeX files found in {final_dir}[/bold red]")
            return None
        return tex_files[0]
    except Exception as e:
        console.print(
            f"[bold red]Error finding LaTeX file in provided directory: {e}[/bold red]"
        )
        return None


def prepare_and_compile_latex(latex_file_path):
    """Prepare and compile LaTeX file to PDF"""
    try:
        pdf_path = os.path.splitext(latex_file_path)[0] + ".pdf"
        if os.path.exists(pdf_path):
            try:
                os.remove(pdf_path)
                console.print("[yellow]Removed existing PDF file[/yellow]")
            except Exception as e:
                console.print(
                    f"[bold yellow]Warning: Could not remove existing PDF: {e}[/bold yellow]"
                )
        console.print(
            f"[yellow]Compiling LaTeX file: {os.path.basename(latex_file_path)}[/yellow]"
        )
        success = compile_latex_to_pdf(latex_file_path)
        if os.path.exists(pdf_path):
            console.print(
                f"[bold green]✓[/bold green] PDF successfully generated: [bold cyan]{pdf_path}[/bold cyan]"
            )
            return True
        else:
            console.print(
                "[bold red]PDF was not created. LaTeX compilation failed.[/bold red]"
            )
            return False
    except Exception as e:
        console.print(f"[bold red]Error compiling LaTeX: {e}[/bold red]")
        return False


def main():
    # Argument parser setup
    parser = argparse.ArgumentParser(
        description="PDF Rerendering Tool: rerender a LaTeX PDF from output files."
    )
    parser.add_argument(
        "-d",
        "--dir",
        dest="final_dir",
        help="Path to the directory containing the final .tex file. If provided, skips search.",
        default=None,
    )
    parser.add_argument(
        "-g",
        "--generated",
        dest="generated_dir",
        help='Path to the generated output root directory to search in (default: "generated").',
        default="generated",
    )
    args = parser.parse_args()

    console.print(
        Panel.fit(
            "[bold cyan]PDF Rerendering Tool[/bold cyan]",
            border_style="blue",
            padding=(1, 10),
        )
    )

    # Determine where to find the .tex file
    if args.final_dir:
        console.print(
            f"[green]Using provided final directory: [bold]{args.final_dir}[/bold][/green]"
        )
        latex_file = find_latex_file_in_dir(args.final_dir)
        if not latex_file:
            return
    else:
        # find latest run directory
        run_dir = find_latest_run_directory(output_dir=args.generated_dir)
        if not run_dir:
            return
        console.print(
            f"[green]Found latest run directory: [bold]{run_dir}[/bold][/green]"
        )
        # find the .tex file in run_dir/final
        latex_file = find_latex_file(run_dir)
        if not latex_file:
            return

    console.print(f"[green]Found LaTeX file: [bold]{latex_file}[/bold][/green]")

    # compile the LaTeX file to PDF
    success = prepare_and_compile_latex(latex_file)

    if success:
        console.print(
            Panel.fit(
                "[bold green]PDF Rerendering Complete![/bold green]",
                border_style="green",
                padding=(1, 10),
            )
        )
    else:
        console.print(
            Panel.fit(
                "[bold red]PDF Rerendering Failed![/bold red]",
                border_style="red",
                padding=(1, 10),
            )
        )


if __name__ == "__main__":
    main()