import gradio as gr
import requests
import os
from openai import OpenAI

UPSTAGE_API_KEY = os.getenv("UPSTAGE_API_KEY")


def extract_text_from_image(filename):
    # Define the API endpoint
    url = "https://api.upstage.ai/v1/document-digitization"

    # Set the authorization header with your API key
    headers = {'Authorization': f'Bearer {UPSTAGE_API_KEY}'}

    # Open the PDF file in binary mode and attach it to the request
    filename = 'letter.jpeg'
    files = {"document": open(filename, "rb")}
    data = {"model": "ocr"}

    response = requests.post(url, headers=headers, files=files, data=data)
        
    if response.status_code == 200:
        text = response.json().get("text", "")
        return text.strip()
    else:
        return f"OCR 실패: {response.status_code} - {response.text}"
        
def translate_text_with_solar(english_text):
    # Initialize the OpenAI client for Solar LLM
    client = OpenAI(
        api_key=UPSTAGE_API_KEY,
        base_url="https://api.upstage.ai/v1"
    )
    prompt = "아래는 영어 편지를 OCR을 통해 추출한 내용입니다. 영어를 한국어로 번역하고, 한국어 어순, 어법에 맞도록 친근한 어투로 윤문한 최종 결과만 한글로 작성해줘. "
    response = client.chat.completions.create(
        model="solar-pro",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.5,
        max_tokens=1024
    )
    return response.choices[0].message.content

with gr.Blocks(title="💌 손글씨 편지 번역기") as demo:
    gr.Markdown("##💌 손글씨 편지 번역기")
    gr.Markdown("📷 편지 이미지를 업로드하면 Upstage Docuemnt OCR이 영어 텍스트를 추출하고,\n🌐 번역하기 버튼을 누르면 Solar LLM을 호출하여 한국어로 번역합니다!")

    with gr.Row():
        # 왼쪽: 이미지 업로드
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="🖼️ 편지 이미지 업로드")

        # 오른쪽: 텍스트 결과
        with gr.Column(scale=2):
            english_box = gr.Textbox(label="📝 추출된 영어 텍스트", lines=10)
            translate_button = gr.Button("🌐 번역하기")
            korean_box = gr.Textbox(label="🇰🇷 번역된 한국어 텍스트", lines=10)
    # with gr.Row():
    #     image_input = gr.Image(type="pil", label="1️⃣ 영어 편지 이미지 업로드")
    #     english_box = gr.Textbox(label="📝 추출된 영어 텍스트", lines=10)

    # with gr.Row():
    #     translate_button = gr.Button("🌐 번역하기")
    #     korean_box = gr.Textbox(label="🇰🇷 번역된 한국어 텍스트", lines=10)

    # Step 1: 이미지 업로드 시 OCR 실행
    image_input.change(fn=extract_text_from_image, inputs=image_input, outputs=english_box)

    # Step 2: 버튼 누르면 번역
    translate_button.click(fn=translate_text_with_solar, inputs=english_box, outputs=korean_box)

if __name__ == "__main__":
    demo.launch()