import gradio as gr
from kiwipiepy import Kiwi
from transformers import pipeline

# 형태소 분석기
kiwi = Kiwi()

# 번역기 설정
translators = {
    "en": pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en"),
    "zh": pipeline("translation", model="Helsinki-NLP/opus-mt-ko-zh"),
    "vi": pipeline("translation", model="Helsinki-NLP/opus-mt-ko-vi")
}

# 문장 나누기
def smart_split(text):
    text = text.replace("다.", "다.★").replace("요.", "요.★").replace("니다.", "니다.★")
    return [s.strip() for s in text.split("★") if s.strip()]

# 핵심 요약 생성
def disaster_style_summary(text):
    parts = smart_split(text)
    keywords = []

    for part in parts:
        tokens = kiwi.tokenize(part)
        nouns = [t.form for t in tokens if t.tag.startswith("N") and len(t.form) > 1]
        verbs = [t.form for t in tokens if t.tag.startswith("V") and len(t.form) > 1]
        if any(v in part for v in ["바랍니다", "주의", "자제", "대피", "착용"]):
            if nouns:
                keywords.append(nouns[0] + " " + (verbs[0] if verbs else "주의"))
        elif nouns:
            keywords.append(nouns[0])
    seen = []
    clean = [k for k in keywords if not (k in seen or seen.append(k))]
    return ". ".join(clean[:5])

# 번역
def translate(text, lang):
    if lang == "ko":
        return text
    try:
        result = translators[lang](text, max_length=100)
        return result[0]["translation_text"]
    except Exception as e:
        return f"[번역 오류] {str(e)}"

# 요약 + 번역 함수
def summarize_and_translate(text, lang):
    if not text.strip():
        return "입력된 내용이 없습니다.", ""
    summary = disaster_style_summary(text)
    translated = translate(summary, lang)
    return summary, translated

# 인터페이스
iface = gr.Interface(
    fn=summarize_and_translate,
    inputs=[
        gr.Textbox(lines=5, label="재난 문자 입력", placeholder="재난안전문자를 여기에 붙여넣으세요."),
        gr.Dropdown(["ko", "en", "zh", "vi"], value="ko", label="번역 언어 선택")
    ],
    outputs=[
        gr.Textbox(label="📌 핵심 요약"),
        gr.Textbox(label="🌍 번역 요약")
    ],
    title="📢 DISASTER SNAP v4"
)

iface.launch()