import os
import datetime
import base64
import tempfile
from io import BytesIO
import uuid

import pytz
import yaml
from PIL import Image
from google import genai
from google.genai import types

from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    tool,
    LiteLLMModel,
)
from smolagents.agent_types import AgentImage
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# === 初始設定 ===
static_tmp_path = tempfile.gettempdir()
os.makedirs(static_tmp_path, exist_ok=True)

# 若部署在 Hugging Face Space，將自動提供 SPACE_HOST
base_url = os.getenv("SPACE_HOST")  # 例如: your-space-id.hf.space

# === 自訂工具 ===

@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """
    A tool that does nothing yet.

    Args:
        arg1: The first argument.
        arg2: The second argument.
    """
    return "What magic will you build?"


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    Fetches the current local time in a specified timezone.

    Args:
        timezone: A string representing a valid timezone (e.g., 'Asia/Taipei').
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


@tool
def image_generation_tool(prompt: str) -> AgentImage:
    """
    Generate an image based on a text prompt using Gemini API.

    Args:
        prompt: The text prompt for image generation.

    Returns:
        AgentImage: An object containing the generated image URL.
    """
    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

    response = client.models.generate_content(
        model="gemini-2.0-flash-exp-image-generation",
        contents=prompt,
        config=types.GenerateContentConfig(
            response_modalities=["TEXT", "IMAGE"]
        ),
    )

    for part in response.candidates[0].content.parts:
        if part.inline_data is not None:
            image = Image.open(BytesIO(part.inline_data.data))
            filename = f"{uuid.uuid4().hex}.png"
            image_path = os.path.join(static_tmp_path, filename)
            image.save(image_path)
            image_url = f"https://{base_url}/file/static/images/{filename}"
            return AgentImage(image_url)

    return "圖片生成失敗"


# === 初始化模型與工具 ===
model = LiteLLMModel(
    model_id="gemini/gemini-2.0-flash-exp",
    api_key=os.getenv("GEMINI_API_KEY"),
)

web_search = DuckDuckGoSearchTool()
final_answer = FinalAnswerTool()

# === 載入 YAML Prompt 模板 ===
with open("prompts.yaml", "r") as stream:
    prompt_templates = yaml.safe_load(stream)

# === 建立代理人 ===
agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        web_search,
        image_generation_tool,
        get_current_time_in_timezone,
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
)

# === 啟動 Gradio 介面 ===
GradioUI(agent).launch()