import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
from googletrans import Translator

# 모델 로드
model_name = "Aoi785/kobart-disaster-summary-v1"
tokenizer = AutoTokenizer.from_pretrained("digit82/kobart-summarization")
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)

# 번역기
translator = Translator()

# 요약 + 번역 함수
def summarize_and_translate(text, lang):
    if not text.strip():
        return "입력된 텍스트가 없습니다.", ""

    # 요약
    try:
        summary = summarizer(text, max_length=50, min_length=10, do_sample=False)[0]['summary_text']
    except Exception as e:
        summary = f"요약 실패: {str(e)}"
        return summary, ""

    # 번역
    try:
        if lang == "ko":
            translated = summary
        else:
            translated = translator.translate(summary, dest=lang).text
    except Exception as e:
        translated = f"번역 실패: {str(e)}"

    return summary, translated

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🆘 재난 문자 요약 및 번역기")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="재난 문자 입력", lines=5, placeholder="예: 내일까지 장마전선 영향으로 많은 비 예상...")
            lang = gr.Dropdown(
                choices=["ko(한국어)", "en(English)", "zh(중국어)", "vi(베트남어)"],
                value="ko(한국어)",
                label="번역 언어 선택"
            )
            run_button = gr.Button("요약 및 번역 실행")

        with gr.Column():
            output_summary = gr.Textbox(label="✅ 요약 결과")
            output_trans = gr.Textbox(label="🌐 번역 결과")

    def process_lang_code(lang_dropdown):
        return lang_dropdown.split("(")[0]

    run_button.click(
        fn=lambda txt, ln: summarize_and_translate(txt, process_lang_code(ln)),
        inputs=[input_text, lang],
        outputs=[output_summary, output_trans]
    )

demo.launch()