import os
import gradio as gr
from inference import run_ocr


def predict(pdf_file, split_page, use_llm, gemini_key):
    if pdf_file is None:
        return "Please upload a PDF.", None
    key = gemini_key or os.getenv("GEMINI_API_KEY", None)
    text, zip_path = run_ocr(pdf_file.name, split_page_enabled=split_page, use_llm=use_llm, gemini_key=key)
    return text, zip_path


with gr.Blocks() as demo:
    gr.Markdown("## PDF OCR (Detectron2 + TrOCR)")
    with gr.Row():
        with gr.Column():
            pdf = gr.File(label="Upload PDF", file_types=[".pdf"])
            split_page = gr.Checkbox(label="Split-page mode", value=False)
            use_llm = gr.Checkbox(label="Gemini correction", value=False)
            gemini_key = gr.Textbox(label="Gemini API Key (optional)", type="password")
            btn = gr.Button("Process")
        with gr.Column():
            out_text = gr.Textbox(label="Extracted Text", lines=18)
            out_zip = gr.File(label="Per-page JSON (ZIP)")
    btn.click(predict, inputs=[pdf, split_page, use_llm, gemini_key], outputs=[out_text, out_zip])


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)