app.py

import gradio as gr
import json
import random
import tempfile
from datetime import datetime
from typing import Dict, Tuple, Optional
from PIL import Image
import io

from google import genai
from google.genai import types
import os
import gtts

API_KEY = os.environ.get("GOOGLE_API_KEY")
client = genai.Client(api_key=API_KEY)

# Pet state storage
pet_state = {
    "name": "",
    "species": "",
    "personality": "",
    "appearance": "",
    "mood": "快樂", # Mood is stored in Chinese, mapped to English for image generation
    "hunger": 80,
    "energy": 90,
    "experience": 0,
    "level": 1,
    "interaction_history": [],
    "battle_stats": {"wins": 0, "losses": 0}
}

def generate_pet_description(custom_input: str = "") -> Dict:
    """Generate pet basic settings using Gemini 2.0 Flash"""
    try:
        if custom_input.strip():
            prompt = f"""
請根據使用者描述創建一隻寶可夢風格的寵物：{custom_input}

請以 JSON 格式返回：
{{
    "name": "寵物名稱",
    "species": "物種類型",
    "personality": "性格描述",
    "appearance": "外觀特徵描述"
}}
"""
        else:
            prompt = """
請隨機創建一隻可愛的寶可夢風格寵物。

請以 JSON 格式返回：
{
    "name": "寵物名稱",
    "species": "物種類型",
    "personality": "性格描述",
    "appearance": "外觀特徵描述"
}

要求：名稱要可愛，性格要有趣，外觀要生動形象。
"""

        response = client.models.generate_content(
            model="gemini-2.5-flash-preview-04-17",
            contents=prompt
        )

        response_text = response.text.strip()
        if response_text.startswith('```json'):
            response_text = response_text[7:-3]
        elif response_text.startswith('```'):
            response_text = response_text[3:-3]

        pet_data = json.loads(response_text)
        return pet_data

    except Exception as e:
        print(f"Error generating pet description: {e}")
        return {
            "name": "皮卡丘",
            "species": "電氣鼠寶可夢",
            "personality": "活潑好動，喜歡交朋友",
            "appearance": "黃色毛髮，紅色臉頰，長長尾巴"
        }

def generate_pet_image(description: str, mood: str = "快樂", context_desc: str = "") -> Optional[Image.Image]:
    """Generate pet image using Gemini 2.0 Flash, including context"""
    try:
        mood_map = {
            "快樂": "happy and cheerful",
            "飢餓": "hungry and tired",
            "興奮": "excited and energetic",
            "疲憊": "sleepy and tired",
            "生氣": "angry and upset",
            "滿足": "satisfied and content"
        }

        mood_desc = mood_map.get(mood, "happy")

        full_description = f"{description}"
        if context_desc:
             full_description += f", {context_desc}"

        image_prompt = f"""
請繪製一張可愛的寶可夢風格寵物插圖，細節如下：
- 描述： {full_description}
- 情緒： {mood_desc}
- 風格： 色彩鮮豔、可愛友善的動漫風格，呈現寵物全身
- 背景： 簡單、柔和的色彩。優先考慮透明背景，若不可行則使用簡單背景
- 品質： 高解析度、細節豐富
- 長寬比：正方形 (1:1)
- **請勿**在圖片中加入任何文字。
"""

        print(f"Generating image with prompt: {image_prompt}")

        image_response = client.models.generate_content(
            model="gemini-2.0-flash-preview-image-generation",
            # model="imagen-3.0-generate-002",
            contents=image_prompt,
            config=types.GenerateContentConfig(
                response_modalities=['TEXT', 'IMAGE']
            )
        )

        if image_response and image_response.candidates and image_response.candidates[0].content.parts:
            for part in image_response.candidates[0].content.parts:
                if part.inline_data and part.inline_data.mime_type.startswith('image/'):
                    image_data_base64 = part.inline_data.data
                    img = Image.open(io.BytesIO(image_data_base64))

                    img = img.resize((512, 512), Image.Resampling.LANCZOS)
                    print("Image generated successfully.")
                    return img

            print("Image data not found in any response parts or not an image.")
            return None
        else:
            print("Image generation response is empty or malformed.")
            return None

    except Exception as e:
        print(f"Error generating pet image: {e}")
        return None

def generate_tts_audio(text: str, lang: str = "zh-TW") -> Optional[str]:
    """Generate text-to-speech audio file"""
    try:
        tts = gtts.gTTS(text=text, lang=lang, slow=False)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        tts.save(temp_file.name)
        return temp_file.name
    except Exception as e:
        print(f"TTS generation failed: {e}")
        return None

def chat_with_pet(user_message: str) -> Tuple[str, Optional[str], None, Optional[Image.Image]]:
    """Chat with the pet"""
    global pet_state

    try:
        if not pet_state["name"]:
             return "請先創建你的寵物！", None, None, None

        context = f"""
你是一隻名叫 {pet_state['name']} 的 {pet_state['species']}，擁有 {pet_state['personality']} 的獨特個性。
你現在的心情是 {pet_state['mood']}，感覺 {pet_state['hunger']}/100 飢餓，精力則有 {pet_state['energy']}/100。

請你完全融入這個角色，以最符合你當前狀態（心情、飢餓、精力）和個性的可愛、友善口吻，回應主人對你說的話："{user_message}"。

回應要求：
1. 保持角色設定並融入當前狀態 (心情、飢餓、精力)。
2. 以可愛、友善的寵物口吻回應。
3. 回應長度約在 20 至 50 字之間。
4. 可以表達你的情感和需求。
5. 使用繁體中文。
"""

        response = client.models.generate_content(
            model="gemini-2.5-flash-preview-04-17",
            contents=context
        )

        pet_response = response.text.strip()

        pet_state["interaction_history"].append({
            "timestamp": datetime.now().strftime("%H:%M:%S"),
            "user": user_message,
            "pet": pet_response,
            "type": "chat"
        })

        audio_file = generate_tts_audio(pet_response)

        image_description = f"{pet_state['species']} - {pet_state.get('appearance', '')}" if pet_state['appearance'] else pet_state['species']
        new_image = generate_pet_image(
            image_description,
            pet_state["mood"],
            context_desc=f"chatting with the user, response: {pet_response}"
        )

        return pet_response, audio_file, None, new_image

    except Exception as e:
        error_msg = f"Chat function error: {e}"
        print(error_msg)
        return "嗚... 我現在說不出話來...", None, None, None

def feed_pet(food_type: str) -> Tuple[str, str, Optional[str], Optional[Image.Image]]:
    """Feed the pet"""
    global pet_state

    try:
        if not pet_state["name"]:
             return "請先創建你的寵物！", "請先創建你的寵物！", None, None

        if not food_type:
            foods = ["樹果", "寶可夢食物", "蘋果", "餅乾", "糖果"]
            food_type = random.choice(foods)

        hunger_increase = random.randint(15, 25)
        pet_state["hunger"] = min(100, pet_state["hunger"] + hunger_increase)

        if pet_state["hunger"] > 80:
            pet_state["mood"] = "快樂"
        elif pet_state["hunger"] > 50:
            pet_state["mood"] = "滿足"

        context = f"""
{pet_state['name']} 剛剛吃了 {food_type}，飢餓程度從之前提升到 {pet_state['hunger']}/100。
請以寵物的角度描述享用食物的感受和反應。

要求：
1. 50-80字的可愛回應
2. 表達對食物的喜愛
3. 用繁體中文
4. 符合 {pet_state['personality']} 的性格
"""

        response = client.models.generate_content(
            model="gemini-2.5-flash-preview-04-17",
            contents=context
        )

        pet_response = response.text.strip()

        pet_state["interaction_history"].append({
            "timestamp": datetime.now().strftime("%H:%M:%S"),
            "action": f"餵食 {food_type}",
            "result": pet_response,
            "type": "feed"
        })

        image_description = f"{pet_state['species']} - {pet_state.get('appearance', '')}" if pet_state['appearance'] else pet_state['species']
        new_image = generate_pet_image(
            image_description,
            pet_state["mood"],
            context_desc=f"eating {food_type}, reaction: {pet_response}"
        )

        status_update = f"等級: {pet_state['level']} | 飢餓: {pet_state['hunger']}/100 | 精力: {pet_state['energy']}/100 | 心情: {pet_state['mood']}"

        audio_file = generate_tts_audio(pet_response)

        return pet_response, status_update, audio_file, new_image

    except Exception as e:
        error_msg = f"Feed function error: {e}"
        print(error_msg)
        status_update = f"等級: {pet_state['level']} | 飢餓: {pet_state['hunger']}/100 | 精力: {pet_state['energy']}/100 | 心情: {pet_state['mood']}"
        return f"嗚... 餵食時出了點問題。({error_msg})", status_update, None, None

def adventure_with_pet() -> Tuple[str, str, Optional[str], Optional[Image.Image]]:
    """Go on an adventure with the pet"""
    global pet_state

    try:
        if not pet_state["name"]:
             return "請先創建你的寵物！", "請先創建你的寵物！", None, None

        energy_cost = random.randint(10, 20)
        pet_state["energy"] = max(0, pet_state["energy"] - energy_cost)

        exp_gain = random.randint(5, 15)
        pet_state["experience"] += exp_gain

        level_up_msg = ""
        while pet_state["experience"] >= pet_state["level"] * 100:
             pet_state["experience"] -= pet_state["level"] * 100
             pet_state["level"] += 1
             level_up_msg += f"恭喜！{pet_state['name']} 升到了 {pet_state['level']} 級！\n"
        level_up_msg = level_up_msg.strip()

        context = f"""
{pet_state['name']} 正在進行一次探險冒險。請創造一個有趣的探險場景：

寵物資訊：
- 種族：{pet_state['species']}
- 性格：{pet_state['personality']}
- 等級：{pet_state['level']}
- 目前精力：{pet_state['energy']}/100

請描述一個短小精彩的探險故事（80-100字），包含：
1. 探險場景
2. 遇到的事件或發現
3. 寵物的反應和收獲
4. 用繁體中文撰寫
"""

        response = client.models.generate_content(
            model="gemini-2.5-flash-preview-04-17",
            contents=context
        )

        adventure_story = response.text.strip()

        image_description = f"{pet_state['species']} - {pet_state.get('appearance', '')}" if pet_state['appearance'] else pet_state['species']
        new_image = generate_pet_image(
            image_description,
            pet_state["mood"],
            context_desc=f"on an adventure: {adventure_story}"
        )

        if level_up_msg:
            adventure_story += f"\n\n🎉 {level_up_msg}"
        adventure_story += f"\n\n📊 獲得經驗值: +{exp_gain}"

        record_result = adventure_story
        pet_state["interaction_history"].append({
            "timestamp": datetime.now().strftime("%H:%M:%S"),
            "action": "探險",
            "result": record_result,
            "type": "adventure",
            "exp_gained": exp_gain
        })

        if pet_state["energy"] > 60:
            pet_state["mood"] = "興奮"
        elif pet_state["energy"] > 30:
            pet_state["mood"] = "滿足"
        else:
            pet_state["mood"] = "疲憊"

        audio_file = generate_tts_audio(adventure_story)

        status_update = f"等級: {pet_state['level']} | 經驗: {pet_state['experience']}/{pet_state['level']*100} | 精力: {pet_state['energy']}/100 | 心情: {pet_state['mood']}"

        return adventure_story, status_update, audio_file, new_image

    except Exception as e:
        error_msg = f"Adventure function error: {e}"
        print(error_msg)
        status_update = f"等級: {pet_state['level']} | 經驗: {pet_state['experience']}/{pet_state['level']*100} | 精力: {pet_state['energy']}/100 | 心情: {pet_state['mood']}"
        return "嗚... 探險時出了點問題。", status_update, None, None

def battle_with_pet() -> Tuple[str, str, Optional[str], Optional[Image.Image]]:
    """Battle with the pet"""
    global pet_state

    try:
        if not pet_state["name"]:
             return "請先創建你的寵物！", "請先創建你的寵物！", None, None

        opponents = ["野生皮卡丘", "小火龍", "傑尼龜", "妙蛙種子", "小拉達", "波波"]
        opponent = random.choice(opponents)

        pet_power = pet_state["level"] * 10 + random.randint(1, 20)
        opponent_power = random.randint(10, 50)

        win = pet_power > opponent_power

        if win:
            pet_state["battle_stats"]["wins"] += 1
            result = "勝利"
            exp_gain = random.randint(10, 20)
            pet_state["experience"] += exp_gain
            pet_state["mood"] = "興奮"
        else:
            pet_state["battle_stats"]["losses"] += 1
            result = "失敗"
            exp_gain = random.randint(3, 8)
            pet_state["experience"] += exp_gain
            pet_state["mood"] = "疲憊"

        level_up_msg = ""
        while pet_state["experience"] >= pet_state["level"] * 100:
             pet_state["experience"] -= pet_state["level"] * 100
             pet_state["level"] += 1
             level_up_msg += f"恭喜！{pet_state['name']} 升到了 {pet_state['level']} 級！\n"
        level_up_msg = level_up_msg.strip()

        energy_cost = random.randint(15, 25)
        pet_state["energy"] = max(0, pet_state["energy"] - energy_cost)

        context = f"""
{pet_state['name']} 剛剛和 {opponent} 進行了一場對戰，結果是{result}。

請描述這場戰鬥的過程和結果（80-120字）：
1. 戰鬥場面描述
2. 使用的招式 (寶可夢風格)
3. 戰鬥結果
4. {pet_state['name']} 的反應
5. 用繁體中文撰寫
6. 符合 {pet_state['personality']} 的性格
"""

        response = client.models.generate_content(
            model="gemini-2.5-flash-preview-04-17",
            contents=context
        )

        battle_story = response.text.strip()

        image_description = f"{pet_state['species']} - {pet_state.get('appearance', '')}" if pet_state['appearance'] else pet_state['species']
        new_image = generate_pet_image(
            image_description,
            pet_state["mood"],
            context_desc=f"in a battle against {opponent}, based on the story: {battle_story}"
        )

        battle_story += f"\n\n📊 獲得經驗值: +{exp_gain}"
        if level_up_msg:
            battle_story += f"\n\n🎉 {level_up_msg}"

        record_result = f"{result} - {battle_story}"
        pet_state["interaction_history"].append({
            "timestamp": datetime.now().strftime("%H:%M:%S"),
            "action": f"對戰 vs {opponent}",
            "result": record_result,
            "type": "battle"
        })

        audio_file = generate_tts_audio(battle_story)

        status_update = f"等級: {pet_state['level']} | 經驗: {pet_state['experience']}/{pet_state['level']*100} | 勝場: {pet_state['battle_stats']['wins']} | 敗場: {pet_state['battle_stats']['losses']} | 精力: {pet_state['energy']}/100 | 心情: {pet_state['mood']}"

        return battle_story, status_update, audio_file, new_image

    except Exception as e:
        error_msg = f"Battle function error: {e}"
        print(error_msg)
        status_update = f"等級: {pet_state['level']} | 經驗: {pet_state['experience']}/{pet_state['level']*100} | 勝場: {pet_state['battle_stats']['wins']} | 敗場: {pet_state['battle_stats']['losses']} | 精力: {pet_state['energy']}/100 | 心情: {pet_state['mood']}"
        return f"嗚... 對戰時出了點問題。", status_update, None, None

def initialize_pet(custom_input: str) -> Tuple[Optional[Image.Image], str, str, str, Optional[str]]:
    """Initialize the pet"""
    global pet_state

    pet_data = generate_pet_description(custom_input)

    pet_state.update({
        "name": pet_data.get("name", "未知寵物"),
        "species": pet_data.get("species", "未知物種"),
        "personality": pet_data.get("personality", "溫和"),
        "appearance": pet_data.get("appearance", ""),
        "mood": "快樂",
        "hunger": 80,
        "energy": 90,
        "experience": 0,
        "level": 1,
        "interaction_history": [],
        "battle_stats": {"wins": 0, "losses": 0}
    })

    image_description = f"{pet_state['species']} - {pet_state['appearance']}" if pet_state['appearance'] else pet_state['species']
    pet_image = generate_pet_image(image_description, pet_state["mood"])

    welcome_msg = f"🎉 歡迎來到寶可夢世界！\n\n" \
                  f"✨ 名稱：{pet_state['name']}\n" \
                  f"🐾 種族：{pet_state['species']}\n" \
                  f"💭 性格：{pet_state['personality']}\n" \
                  f"😊 心情：{pet_state['mood']}\n\n" \
                  f"現在你可以和 {pet_state['name']} 互動了！"

    status_info = f"等級: {pet_state['level']} | 飢餓: {pet_state['hunger']}/100 | 精力: {pet_state['energy']}/100 | 心情: {pet_state['mood']}"

    audio_file = generate_tts_audio(f"你好！我是{pet_state['name']}，很高興認識你！")

    return pet_image, welcome_msg, status_info, "", audio_file

def get_interaction_history() -> str:
    """Get interaction history"""
    if not pet_state["interaction_history"]:
        return "還沒有互動記錄呢！快和你的寵物互動吧！"

    history_text = f"📋 {pet_state['name']} 的互動日誌\n" + "="*50 + "\n\n"

    for record in pet_state["interaction_history"][-20:]:
        time_stamp = record["timestamp"]
        if record["type"] == "chat":
            history_text += f"🕐 {time_stamp} [聊天]\n"
            history_text += f"👤 你: {record['user']}\n"
            history_text += f"🐾 {pet_state['name']}: {record['pet']}\n\n"
        else:
            action_desc = record.get('action', record['type'])
            result_desc = record.get('result', '無結果描述')
            history_text += f"🕐 {time_stamp} [{action_desc}]\n"
            history_text += f"📝 {result_desc}\n\n"

    return history_text

def create_app():
    """Create the Gradio application"""

    with gr.Blocks(
        title="Virtual Pokemon Companion",
        theme=gr.themes.Soft(),
        css="""
        .main-container { max-width: 1200px; margin: 0 auto; }
        .pet-image { border-radius: 15px; box-shadow: 0 4px 8px rgba(0,0,0,0.1); object-fit: contain; }
        .status-box { background: linear-gradient(45deg, #ff9a9e, #fecfef); padding: 15px; border-radius: 10px; }
        .interaction-box { background: linear-gradient(45deg, #a8edea, #fed6e3); padding: 15px; border-radius: 10px; }
        """) as app:

        gr.Markdown("""
        # 🌟 Virtual Pokemon Companion
        ### An AI-powered pet system using Google Gemini
        """)

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## 🐾 Your Pet")

                pet_image = gr.Image(
                    label="Pet Image",
                    height=300,
                    show_label=False,
                    elem_classes=["pet-image"],
                    interactive=False
                )

                pet_status = gr.Textbox(
                    label="Pet Status",
                    value="Please create your pet first!",
                    interactive=False,
                    lines=3,
                    elem_classes=["status-box"]
                )

                gr.Markdown("### 🎮 Create Pet")
                custom_pet_input = gr.Textbox(
                    label="Custom Pet Traits (Optional)",
                    placeholder="e.g., A flying blue dragon with a gentle personality..."
                )

                with gr.Row():
                    create_random_btn = gr.Button("🎲 Randomize", variant="primary")
                    create_custom_btn = gr.Button("✨ Create Custom", variant="secondary")

            with gr.Column(scale=1):
                gr.Markdown("## 💬 Interaction Area")

                chat_output = gr.Textbox(
                    label="Pet Response",
                    value="Your pet is waiting to be created...",
                    interactive=False,
                    lines=4,
                    elem_classes=["interaction-box"]
                )

                audio_output = gr.Audio(
                    label="Pet Voice",
                    visible=True,
                    autoplay=True
                )

                gr.Markdown("### 💭 Chat")
                chat_input = gr.Textbox(
                    label="Talk to your pet",
                    placeholder="Enter your message here..."
                )
                chat_btn = gr.Button("💬 Chat", variant="primary")

                gr.Markdown("### 🎯 Activities")
                with gr.Row():
                    feed_btn = gr.Button("🍎 Feed", variant="secondary")
                    adventure_btn = gr.Button("🗺️ Adventure", variant="secondary")

                with gr.Row():
                    battle_btn = gr.Button("⚔️ Battle", variant="secondary")
                    history_btn = gr.Button("📋 View Log", variant="secondary")

        with gr.Row():
            history_output = gr.Textbox(
                label="📜 Interaction Log",
                lines=8,
                interactive=False,
                visible=False,
                max_lines=20
            )

        # Event bindings
        create_random_btn.click(
            fn=lambda: initialize_pet(""),
            outputs=[pet_image, chat_output, pet_status, custom_pet_input, audio_output]
        )

        create_custom_btn.click(
            fn=initialize_pet,
            inputs=[custom_pet_input],
            outputs=[pet_image, chat_output, pet_status, custom_pet_input, audio_output]
        )

        chat_btn.click(
            fn=chat_with_pet,
            inputs=[chat_input],
            outputs=[chat_output, audio_output, chat_input, pet_image]
        )

        chat_input.submit(
            fn=chat_with_pet,
            inputs=[chat_input],
            outputs=[chat_output, audio_output, chat_input, pet_image]
        )

        feed_btn.click(
            fn=lambda: feed_pet(""),
            outputs=[chat_output, pet_status, audio_output, pet_image]
        )

        adventure_btn.click(
            fn=adventure_with_pet,
            outputs=[chat_output, pet_status, audio_output, pet_image]
        )

        battle_btn.click(
            fn=battle_with_pet,
            outputs=[chat_output, pet_status, audio_output, pet_image]
        )

        def toggle_history():
            current_visibility = history_output.visible
            new_visibility = not current_visibility
            history_content = get_interaction_history() if new_visibility else ""
            return gr.update(visible=new_visibility), history_content

        history_btn.click(
            fn=toggle_history,
            outputs=[history_output, history_output]
        )

    return app

app = create_app()

app.launch()