import json
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from gtts import gTTS
from diffusers import StableDiffusionPipeline
import gradio as gr

def load_fairytale(file_obj):
    data = json.loads(file_obj.read().decode("utf-8"))  # read & decode
    return data['title'], data['content']

def generate_grandma_voice(text):
    grandma_text = f"에구구 얘야, 잘 들어보렴. {text.strip()} ... 옛날 옛적 이야기란다~"
    tts = gTTS(text=grandma_text, lang='ko')
    audio_path = "grandma_voice.mp3"
    tts.save(audio_path)
    return audio_path

emotion_tokenizer = AutoTokenizer.from_pretrained("monologg/koelectra-base-discriminator")
emotion_model = AutoModelForSequenceClassification.from_pretrained("monologg/koelectra-base-discriminator")

def classify_emotion(text):
    inputs = emotion_tokenizer(text, return_tensors="pt", truncation=True)
    with torch.no_grad():
        outputs = emotion_model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=1)
    label = torch.argmax(probs).item()
    emotions_ko = ["기쁨", "슬픔", "분노", "불안", "중립"]
    emotions_en = ["joy", "sadness", "anger", "anxiety", "neutral"]
    return emotions_en[label], emotions_ko[label]

stable_pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16
)
device = "cuda" if torch.cuda.is_available() else "cpu"
stable_pipe = stable_pipe.to(device)

def generate_emotion_image(emotion_en):
    prompt = f"A dreamy digital painting that represents the feeling of {emotion_en}"
    image = stable_pipe(prompt).images[0]
    image_path = f"{emotion_en}_image.png"
    image.save(image_path)
    return image_path

def run_all(fairytale_file, child_feeling_text):
    title, content = load_fairytale(fairytale_file)
    audio_path = generate_grandma_voice(content[:300])
    emotion_en, emotion_ko = classify_emotion(child_feeling_text)
    image_path = generate_emotion_image(emotion_en)
    return title, audio_path, emotion_ko, image_path

demo = gr.Interface(
    fn=run_all,
    inputs=[
        gr.File(label="동화 JSON 파일 업로드"),
        gr.Textbox(label="아이의 감상문")
    ],
    outputs=[
        gr.Text(label="동화 제목"),
        gr.Audio(label="할머니 목소리"),
        gr.Text(label="감정 분석 결과 (한국어)"),
        gr.Image(label="감정 표현 이미지")
    ],
    title="AI 할머니가 읽어주는 감성 동화책",
    description="동화를 업로드하면 할머니가 읽어주고, 아이 감상문에 맞춰 감정 이미지를 생성합니다."
)

if __name__ == "__main__":
    demo.launch()