import gradio as gr
import os
from pathlib import Path
from redact_extract import make_side_by_side, make_overlay_white, compute_redaction_stats

def process_pdf(pdf_file, mode, line_tol, space_unit, min_spaces, match_font):
    if pdf_file is None:
        return None, "Please upload a PDF file."

    input_pdf_path = pdf_file.name
    
    # Create a unique output path in a temporary directory
    output_dir = Path("unredacted_output")
    output_dir.mkdir(exist_ok=True)
    
    pdf_name = Path(input_pdf_path).stem
    suffix = f"_{mode}.pdf"
    output_pdf_path = output_dir / (pdf_name + suffix)

    if mode == "side_by_side":
        make_side_by_side(
            input_pdf_path, 
            output_pdf_path, 
            line_tol=line_tol, 
            space_unit_pts=space_unit, 
            min_spaces=min_spaces, 
            match_font=match_font
        )
    else: # overlay_white
        make_overlay_white(
            input_pdf_path, 
            output_pdf_path, 
            line_tol=line_tol, 
            space_unit_pts=space_unit, 
            min_spaces=min_spaces, 
            match_font=match_font
        )

    stats = compute_redaction_stats(input_pdf_path, line_tol=line_tol)
    stats_display = stats.display()

    return str(output_pdf_path), stats_display

# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Unredact PDF")
    gr.Markdown("Upload a PDF with redactions to recover the text underneath.")

    with gr.Row():
        with gr.Column(scale=1):
            pdf_input = gr.File(label="Input PDF", file_types=[".pdf"])
            
            gr.Markdown("### Processing Options")
            mode = gr.Radio(["side_by_side", "overlay_white"], label="Output Mode", value="side_by_side")
            line_tol = gr.Slider(1.0, 5.0, value=2.0, label="Line Tolerance", info="Line grouping tolerance (pts). Try 1.5–4.0")
            space_unit = gr.Slider(1.0, 10.0, value=3.0, label="Space Unit", info="Points per inserted space (bigger => fewer spaces)")
            min_spaces = gr.Slider(0, 5, value=1, step=1, label="Minimum Spaces")
            match_font = gr.Checkbox(label="Attempt to match font", value=False)
            
            process_button = gr.Button("Process PDF", variant="primary")

        with gr.Column(scale=2):
            gr.Markdown("### Output")
            pdf_output = gr.File(label="Processed PDF")
            stats_output = gr.Textbox(label="Redaction Statistics", lines=15)

    process_button.click(
        fn=process_pdf,
        inputs=[pdf_input, mode, line_tol, space_unit, min_spaces, match_font],
        outputs=[pdf_output, stats_output]
    )
    
    gr.Examples(
        examples=[
            ["examples/an_example.png", "side_by_side", 2.0, 3.0, 1, False],
        ],
        inputs=[pdf_input, mode, line_tol, space_unit, min_spaces, match_font],
        outputs=[pdf_output, stats_output],
        fn=process_pdf,
        cache_examples=True
    )

if __name__ == "__main__":
    demo.launch()