import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# 한국어 -> 영어 모델 로드
model_name_ko_en = "Helsinki-NLP/opus-mt-ko-en"
tokenizer_ko_en = MarianTokenizer.from_pretrained(model_name_ko_en)
model_ko_en = MarianMTModel.from_pretrained(model_name_ko_en)

# 영어 -> 한국어 모델 로드
model_name_en_ko = "Helsinki-NLP/opus-mt-en-ko"
tokenizer_en_ko = MarianTokenizer.from_pretrained(model_name_en_ko)
model_en_ko = MarianMTModel.from_pretrained(model_name_en_ko)

def translate_text(text, source_lang, target_lang):
    """
    텍스트를 번역하는 함수.
    """
    try:
        if source_lang == "한국어" and target_lang == "영어":
            # 한국어 -> 영어 번역
            tokenized = tokenizer_ko_en(text, return_tensors="pt", padding=True)
            translated = model_ko_en.generate(**tokenized)
            result = tokenizer_ko_en.decode(translated[0], skip_special_tokens=True)
        elif source_lang == "영어" and target_lang == "한국어":
            # 영어 -> 한국어 번역
            tokenized = tokenizer_en_ko(text, return_tensors="pt", padding=True)
            translated = model_en_ko.generate(**tokenized)
            result = tokenizer_en_ko.decode(translated[0], skip_special_tokens=True)
        else:
            return "지원하지 않는 언어 조합입니다."
        return result
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"

# Gradio UI 구성
with gr.Blocks() as demo:
    gr.Markdown("# 한국어 ↔ 영어 번역기")

    # 언어 선택
    source_lang = gr.Radio(
        choices=["한국어", "영어"],
        label="원본 언어",
        value="한국어"
    )
    target_lang = gr.Radio(
        choices=["한국어", "영어"],
        label="목표 언어",
        value="영어"
    )

    # 입력 텍스트
    input_text = gr.Textbox(label="번역할 텍스트", lines=5)

    # 번역 결과
    translation_output = gr.Textbox(label="번역 결과", lines=5, interactive=False)

    # 번역 버튼
    translate_button = gr.Button("번역")

    # 번역 함수 연결
    translate_button.click(
        fn=translate_text,
        inputs=[input_text, source_lang, target_lang],
        outputs=translation_output
    )

# 메인 실행부
if __name__ == "__main__":
    demo.launch()