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

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

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):
    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"
    }

    # 해당 부분을 ["text"]로 변경하여 formats를 리스트로 설정
    data = {
        "src": f"data:image/jpeg;base64,{image_base64}",
        "formats": ["text"]
    }

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

    result = response.json()
    # formats_results 변수는 사용하지 않으므로 제거하여 코드를 단순화
    return result
  

def build_gradio_app(css=css):
    with gr.Blocks(css=css) as demo:
        gr.Markdown("# ArXivGPT OCR LaTeX")
        gr.Markdown("변환하기 버튼 누르고 출력된 결과값을 '복사'하여 화면 아래 박스('텍스트 지우고 빈 상태로')에 붙여넣고 'Render HTML'버튼을 누르세요.")
        
        image_input = gr.Image(type="pil", label="이미지 업로드")

        def process_and_output(image):
            # 이미지 처리 및 결과 반환 로직
            latex_result = get_latex_from_image_all_formats(image)
            text_result = latex_result.get("text", "No result")
            return text_result

        submit_button = gr.Button("변환하기")
        outputs = gr.Textbox(label="결과")

        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])
        
        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: 8px;"></iframe>', elem_id="latex_iframe")

    return demo

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