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

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

# Mathpix에 이미지 파일을 보내어 여러 포맷의 LaTeX 문자열을 추출하는 함수
def get_latex_from_image_all_formats(image):
    # PIL 이미지를 바이트로 변환
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

    # Mathpix API 요청 헤더
    headers = {
        "app_id": app_id,
        "app_key": app_key,
        "Content-Type": "application/json"
    }

    # Mathpix API 요청 바디
    data = {
        "src": f"data:image/jpeg;base64,{image_base64}",
        "formats": "text"
    }

    # Mathpix API 요청 보내기
    response = requests.post(app_url, headers=headers, json=data)
    response.raise_for_status()  # 요청 실패 시 예외 발생

    # 응답에서 각 포맷의 LaTeX 추출
    result = response.json()
    formats_results = {f: result.get(f, f"{f} 추출 실패") for f in data["formats"]}

    return formats_results


def build_gradio_app():
    with gr.Blocks() as demo:
        with gr.Row():
            image_input = gr.Image(type="pil", label="이미지 업로드")
            submit_button = gr.Button("변환하기")

        outputs = [gr.Textbox(label=f"{f} 결과") for f in "text"]

        latex_iframe = gr.HTML(value='<iframe src="https://www.mathjax.org/#demo" style="width: 100%; height: 700px; border: 2px solid #007bff; border-radius: 8px;"></iframe>', elem_id="latex_iframe")
        

        def process_and_output(image):
            latex_result = get_latex_from_image_all_formats(image)  # 변수명 수정
            return [latex_result.get(f) for f in "text"]  # 수정된 부분
        
        submit_button.click(fn=process_and_output, inputs=image_input, outputs=outputs)

    return demo

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