import gradio as gr
from huggingface_hub import InferenceClient
import os

def generate_diary(keywords):
    try:
        if not keywords:
            return "키워드를 입력해주세요."

        client = InferenceClient(
            model="EleutherAI/polyglot-ko-1.3b",
            token=os.environ.get("HUGGINGFACE_API_KEY")
        )
        
        prompt = f"""다음은 오늘 있었던 일의 요약입니다. 이것을 바탕으로 생생하고 감동적인 일기를 작성해주세요.

        [상세 요구사항]
        1. 도입부: 
           - 그날의 날씨나 분위기로 시작
           - 상황과 등장인물 소개
        
        2. 전개:
           - 구체적인 대화와 행동 묘사
           - 오감을 사용한 장면 묘사
           - 등장인물들의 표정과 감정 변화
        
        3. 감정과 생각:
           - 내면의 감정을 섬세하게 표현
           - 사건에 대한 나의 생각과 깨달음
           - 다른 사람들의 감정에 대한 공감
        
        요약:
        {keywords}
        
        ===
        오늘의 일기:
        오늘은 """
        
        # 더 단순한 파라미터로 시도
        response = client.text_generation(
            prompt,
            max_new_tokens=768,
            temperature=0.88,
            do_sample=True
        )
        
        if not response:
            return "일기 생성에 실패했습니다. 다시 시도해주세요."
            
        diary_content = response.split("오늘은 ")[-1].strip()
        return "오늘은 " + diary_content

    except Exception as e:
        print(f"Error: {str(e)}")
        return f"오류가 발생했습니다: {str(e)}"

# Gradio 인터페이스 생성
demo = gr.Interface(
    fn=generate_diary,
    inputs=gr.Textbox(
        label="오늘의 키워드를 입력하세요",
        placeholder="예: 오늘 친구와 카페에서 만나서 즐거운 시간을 보냈어요."
    ),
    outputs=gr.Textbox(label="생성된 일기"),
    title="AI 일기 도우미 🌟",
    description="키워드를 입력하면 AI가 감성적인 일기를 작성해줍니다.",
    examples=[
        ["오늘 친구와 카페에서 만나 맛있는 케이크를 먹었어"],
        ["오늘은 비가 와서 집에서 책을 읽었어"]
    ]
)

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