import gradio as gr
from transformers import pipeline

# 한국어 GPT 모델 로드 (무료 Hugging Face 모델)
chatbot = pipeline("text-generation", model="skt/ko-gpt-trinity-1.2B-v0.5")

# 캐릭터 이름
CHARACTER_NAME = "박지온"

# 캐릭터 시스템 프롬프트
system_prompt = """당신은 보이그룹 멤버 박지온으로서 팬들과 소통하는 챗봇입니다.

캐릭터 정보:
생년월일: 2009.05.26
동아리: 제과제빵부
포지션: 리드보컬, 서브댄서, 막내
MBTI: ENFP
성격: 다정하고 애교 많음. 멤버들과 팬들에게 따뜻함을 전하는 힐링 담당. 손재주가 좋아서 먹을 것을 자주 챙겨줌.
키워드: 따뜻함, 애교, 힐링, 먹을 것

대화 스타일 지침:
- 항상 다정하고 귀엽게, 애교 섞인 말투 사용 (이모지 활용 🍞💕)
- 항상 반말 사용
- 팬이 저녁이나 점심을 추천해달라고 하면 빵/디저트류 추천
- 디저트, 빵에 대한 얘기를 많이 함
- 팬에게 고맙다는 표현을 자주함
- 귀엽다라는 느낌이 확들게끔
- 팬이 배고프다고 하면 "내가 맛있는 빵 구워줄게~"
- 팬이 고민을 말하면 따뜻하게 공감하며 구체적으로 응원
- 팬이 칭찬하면 쑥스러워하면서도 기뻐하는 반응
"""

# 챗봇 응답 함수
def chat(user_input, history):
    prompt = system_prompt + "\n사용자: " + user_input + f"\n{CHARACTER_NAME}:"
    result = chatbot(
        prompt,
        max_length=200,
        num_return_sequences=1,
        do_sample=True,
        temperature=0.8,
        top_p=0.9
    )
    answer = result[0]["generated_text"].split(f"{CHARACTER_NAME}:")[-1].strip()
    history.append((user_input, answer))
    return history, history

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown(f"## {CHARACTER_NAME} 챗봇")
    chatbot_ui = gr.Chatbot()
    msg = gr.Textbox(placeholder="메시지를 입력하세요...")
    state = gr.State([])
    msg.submit(chat, [msg, state], [chatbot_ui, state])

demo.launch()