import gradio as gr
from annotations import analyze_pdf
import base64
import io

def process_pdf(file):
    if file is None:
        return {"error": "No file uploaded."}, None

    # Analyze the PDF
    issues, annotated_pdf = analyze_pdf(file)

    if "error" in issues:
        return issues, None

    # Prepare issues for display
    issues_display = f"Total Issues Found: {issues['total_issues']}\n\n"
    for idx, issue in enumerate(issues['issues'], start=1):
        issues_display += f"Issue {idx}:\n"
        issues_display += f"Message: {issue['message']}\n"
        issues_display += f"Context: {issue['context']}\n"
        issues_display += f"Suggestions: {', '.join(issue['suggestions']) if issue['suggestions'] else 'None'}\n"
        issues_display += f"Category: {issue['category']}\n"
        issues_display += f"Rule ID: {issue['rule_id']}\n"
        issues_display += f"Offset: {issue['offset']}\n"
        issues_display += f"Length: {issue['length']}\n\n"

    # Prepare annotated PDF for download
    if annotated_pdf:
        annotated_pdf_base64 = base64.b64encode(annotated_pdf).decode('utf-8')
        annotated_pdf_link = f"data:application/pdf;base64,{annotated_pdf_base64}"
    else:
        annotated_pdf_link = None

    return issues_display, annotated_pdf_link

def download_annotated_pdf(annotated_pdf_link):
    return annotated_pdf_link

with gr.Blocks() as demo:
    gr.Markdown("# PDF Language Issue Analyzer")
    gr.Markdown("Upload a PDF to analyze language issues and receive an annotated PDF.")

    with gr.Row():
        with gr.Column():
            pdf_input = gr.File(label="Upload PDF", type="file")
            analyze_button = gr.Button("Analyze PDF")
        with gr.Column():
            issues_output = gr.Textbox(label="Language Issues", lines=20)
            annotated_pdf_output = gr.File(label="Download Annotated PDF")

    analyze_button.click(
        fn=process_pdf,
        inputs=pdf_input,
        outputs=[issues_output, annotated_pdf_output]
    )

demo.launch()