import gradio as gr
from transformers import pipeline, MarianMTModel, MarianTokenizer

# ✅ 요약 파이프라인 초기화
summarizer = pipeline("summarization", model="digit82/kobart-summarization")

# ✅ 번역 모델 사전 정의
translation_models = {
    "en": "Helsinki-NLP/opus-mt-ko-en",
    "zh": "Helsinki-NLP/opus-mt-ko-zh",
    "vi": "Helsinki-NLP/opus-mt-ko-vi"
}

# ✅ 번역 함수
def translate(text, target_lang):
    if target_lang == "ko":
        return text
    model_name = translation_models[target_lang]
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    inputs = tokenizer.prepare_seq2seq_batch([text], return_tensors="pt", padding=True)
    translated = model.generate(**inputs)
    return tokenizer.decode(translated[0], skip_special_tokens=True)

# ✅ 메인 함수
def summarize_and_translate(text, target_lang):
    text = text[:300]  # 입력 제한
    summary = summarizer(text, max_length=20, min_length=5, do_sample=False)[0]['summary_text']
    translated = translate(summary, target_lang)
    return summary, translated

# ✅ UI 구성
lang_choices = [
    ("ko", "한국어 (ko)"),
    ("en", "English (en)"),
    ("zh", "中国人 (zh)"),
    ("vi", "Tiếng Việt (vi)")
]

with gr.Blocks() as demo:
    gr.Markdown("## 📢 재난문자 초간단 요약 & 번역기")

    with gr.Row():
        textbox = gr.Textbox(lines=5, label="재난 문자 입력", placeholder="여기에 재난 문자를 입력하세요 (최대 300자)")

    with gr.Row():
        lang = gr.Dropdown(
            choices=[code for code, _ in lang_choices],
            label="번역 언어 선택",
            value="ko",
            info="ko(한국어), en(English), zh(중국어), vi(베트남어)"
        )

    with gr.Row():
        summary_out = gr.Textbox(label="✅ 요약 결과", lines=2)
        translate_out = gr.Textbox(label="🌐 번역 결과", lines=2)

    btn = gr.Button("요약 및 번역 실행")
    btn.click(fn=summarize_and_translate, inputs=[textbox, lang], outputs=[summary_out, translate_out])

demo.launch()