import gradio as gr
import torch
import librosa
from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, pipeline

# -------------------------------
# 1) STT 모델 로드 (Whisper 기반)
# -------------------------------
processor = AutoProcessor.from_pretrained("openai/whisper-small")
model = AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-small")

# -------------------------------
# 2) 감정 분석 모델 로드
# -------------------------------
sentiment_pipe = pipeline(
    "sentiment-analysis",
    model="monologg/koelectra-base-v3-discriminator",
    tokenizer="monologg/koelectra-base-v3-discriminator"
)

# -------------------------------
# 3) 오디오 -> 텍스트
# -------------------------------
def transcribe_audio(audio_path):
    speech, _ = librosa.load(audio_path, sr=16000)
    input_features = processor(speech, sampling_rate=16000, return_tensors="pt").input_features
    predicted_ids = model.generate(input_features)
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    return transcription

# -------------------------------
# 4) 감정 레이블 -> 색상 매핑
# -------------------------------
def label_to_color(label):
    if label in ["4 stars", "5 stars", "LABEL_4", "LABEL_5"]:
        return "green"  # 긍정
    elif label in ["1 star", "2 stars", "LABEL_1", "LABEL_2"]:
        return "red"    # 부정
    else:  # 중립
        return "orange"

# -------------------------------
# 5) 텍스트 감정 분석 (문장 전체)
# -------------------------------
def sentiment_whole_text(text):
    res = sentiment_pipe(text)[0]
    label = res['label']
    score = res['score']
    color = label_to_color(label)
    styled_text = f"<span style='color:{color}'>{text}</span>"
    legend = (
        "<div style='margin-top:10px;'>"
        "<b>색상 설명:</b> "
        "<span style='color:green'>녹색=긍정</span>, "
        "<span style='color:red'>빨강=부정</span>, "
        "<span style='color:orange'>주황=중립/보통</span>"
        "</div>"
    )
    return styled_text + legend, f"감정: {label}, 신뢰도: {score:.2f}"

# -------------------------------
# 6) 텍스트 감정 분석 (단어별)
# -------------------------------
def sentiment_word_level(text):
    words = text.split()
    styled_words = []
    for w in words:
        res = sentiment_pipe(w)[0]
        label = res['label']
        color = label_to_color(label)
        styled_words.append(f"<span style='color:{color}'>{w}</span>")
    styled_text = " ".join(styled_words)
    legend = (
        "<div style='margin-top:10px;'>"
        "<b>색상 설명:</b> "
        "<span style='color:green'>녹색=긍정</span>, "
        "<span style='color:red'>빨강=부정</span>, "
        "<span style='color:orange'>주황=중립/보통</span>"
        "</div>"
    )
    return styled_text + legend, "단어별 감정 표시 완료"

# -------------------------------
# 7) 오디오 -> 텍스트 + 감정 분석 (문장+단어)
# -------------------------------
def process_audio_full(audio_file):
    text = transcribe_audio(audio_file)
    whole_text_result, whole_text_score = sentiment_whole_text(text)
    word_level_result, word_level_status = sentiment_word_level(text)
    return text, whole_text_result, whole_text_score, word_level_result, word_level_status

# -------------------------------
# 8) Gradio UI 구성
# -------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🎤 오디오/텍스트 → 감정 분석")
    
    with gr.Tabs():
        # ------------------- 오디오 -> 텍스트 -------------------
        with gr.Tab("오디오 → 텍스트"):
            audio_input_1 = gr.Audio(label="음성 업로드", type="filepath")
            audio_text_output = gr.Textbox(label="변환된 텍스트")
            audio_transcribe_btn = gr.Button("텍스트 추출")
            audio_transcribe_btn.click(fn=transcribe_audio, inputs=[audio_input_1], outputs=[audio_text_output])
        
        # ------------------- 텍스트 -> 감정 분석 -------------------
        with gr.Tab("텍스트 → 감정 분석"):
            text_input = gr.Textbox(label="텍스트 입력")
            sentiment_whole_output = gr.HTML(label="문장 단위 감정 분석")
            sentiment_whole_score = gr.Markdown(label="감정 결과")
            sentiment_word_output = gr.HTML(label="단어 단위 감정 분석")
            sentiment_btn = gr.Button("감정 분석")
            def analyze_text(text):
                whole_res, whole_score = sentiment_whole_text(text)
                word_res, word_status = sentiment_word_level(text)
                return whole_res, whole_score, word_res
            sentiment_btn.click(
                fn=analyze_text,
                inputs=[text_input],
                outputs=[sentiment_whole_output, sentiment_whole_score, sentiment_word_output]
            )
        
        # ------------------- 오디오 → 텍스트 + 감정 분석 -------------------
        with gr.Tab("오디오 → 텍스트 + 감정 분석"):
            audio_input_2 = gr.Audio(label="음성 업로드", type="filepath")
            audio_text_output_2 = gr.Textbox(label="변환된 텍스트")
            sentiment_whole_output_2 = gr.HTML(label="문장 단위 감정 분석")
            sentiment_whole_score_2 = gr.Markdown(label="감정 결과")
            sentiment_word_output_2 = gr.HTML(label="단어 단위 감정 분석")
            audio_process_btn = gr.Button("분석 시작")
            def process_audio_tab(audio_file):
                text, whole_res, whole_score, word_res, word_status = process_audio_full(audio_file)
                return text, whole_res, whole_score, word_res
            audio_process_btn.click(
                fn=process_audio_tab,
                inputs=[audio_input_2],
                outputs=[audio_text_output_2, sentiment_whole_output_2, sentiment_whole_score_2, sentiment_word_output_2]
            )

demo.launch()