import gradio as gr
import requests
import json
import base64
from PIL import Image
import io
import os

css = """
footer { 
    visibility: hidden; 
}

.gradio-container {
    background: linear-gradient(145deg, #f0f0f0, #ffffff);
    box-shadow: 20px 20px 60px #bebebe, -20px -20px 60px #ffffff;
    border-radius: 15px;
}

.gr-button {
    background: linear-gradient(145deg, #007bff, #0056b3);
    border: none;
    padding: 12px 24px;
    border-radius: 10px;
    color: white;
    font-weight: bold;
    box-shadow: 5px 5px 15px rgba(0,0,0,0.2);
    transition: transform 0.2s;
}

.gr-button:hover {
    transform: translateY(-2px);
    box-shadow: 7px 7px 20px rgba(0,0,0,0.25);
}

.gr-input, .gr-box {
    border-radius: 10px;
    box-shadow: inset 5px 5px 10px #d1d1d1, inset -5px -5px 10px #ffffff;
    border: 2px solid #e0e0e0;
}

.gr-form {
    background: #ffffff;
    padding: 20px;
    border-radius: 15px;
    box-shadow: 0 10px 30px rgba(0,0,0,0.1);
}
"""

app_id = os.getenv("app_id")
app_key = os.getenv("app_key")
app_url = os.getenv("app_url")

def get_latex_from_image_all_formats(image):
    try:
        if image is None:
            return {"text": "No image provided"}
            
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

        headers = {
            "app_id": app_id,
            "app_key": app_key,
            "Content-Type": "application/json"
        }

        data = {
            "src": f"data:image/jpeg;base64,{image_base64}",
            "formats": ["text"]
        }

        response = requests.post(app_url, headers=headers, json=data)
        response.raise_for_status()

        return response.json()
    except Exception as e:
        return {"text": f"Error: {str(e)}"}

def build_gradio_app(css=css):
    with gr.Blocks(css=css) as demo:
        gr.Markdown("# ArXivGPT OCR LaTeX Converter")
        gr.Markdown("Upload an image, click 'Convert' button, and copy the output. Paste it in the empty box below and click 'Render HTML'.")
        
        image_input = gr.Image(type="pil", label="Upload Image")

        def process_and_output(image):
            try:
                if image is None:
                    return "Please upload an image"
                    
                latex_result = get_latex_from_image_all_formats(image)
                if latex_result is None:
                    return "Error processing image"
                    
                return latex_result.get("text", "No result")
            except Exception as e:
                return f"Error: {str(e)}"

        submit_button = gr.Button("Convert", elem_classes="custom-button")
        outputs = gr.Textbox(label="Output")

        submit_button.click(fn=process_and_output, inputs=[image_input], outputs=[outputs])

        examples = gr.Examples(
            examples=[["ko.png"], ["en.png"], ["hand.jpg"]], 
            inputs=[image_input], 
            fn=process_and_output, 
            outputs=[outputs],
            label="Example Images"
        )
        
        latex_iframe = gr.HTML(
            value='<iframe src="https://mathjax.github.io/MathJax-demos-web/input-tex_mml2chtml.html" style="width: 100%; height: 700px; border: 2px solid #007bff; border-radius: 12px; box-shadow: 0 10px 30px rgba(0,0,0,0.15);"></iframe>',
            elem_id="latex_iframe"
        )

    return demo

if __name__ == "__main__":
    app = build_gradio_app()
    app.launch()