# kids_museum_agent.py

from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    HfApiModel,
    load_tool,
    tool
)
import datetime
import requests
import pytz
import yaml
import os

from tools.final_answer import FinalAnswerTool  # Adjust import path if needed

########################################
# UTILITY FUNCTION:
# Restructure Europeana JSON response
########################################
def restructure_europeana_response(europeana_json: dict) -> str:
    """
    Takes a Europeana Search API JSON response and returns a
    simplified string summary, focusing on child-friendly information.
    """
    # Check if Europe's response is successful
    if not europeana_json.get("success"):
        return f"Oops! There was a problem with Europeana: {europeana_json.get('error', 'Unknown error')}"

    items = europeana_json.get("items", [])
    if not items:
        return "Sorry, I couldn’t find anything on Europeana for that topic!"

    # Build a short list of results
    result_lines = []
    for idx, item in enumerate(items, start=1):
        # Try to get a title
        title_list = item.get("title") or []
        if not title_list:
            # fallback: check dcTitleLangAware if "title" is missing
            dc_title_lang = item.get("dcTitleLangAware", {})
            if dc_title_lang:
                first_lang = next(iter(dc_title_lang))
                title_list = dc_title_lang[first_lang]
        title_str = title_list[0] if title_list else "[No title found]"

        # Provider (museum/institution)
        provider_list = item.get("dataProvider") or []
        provider_str = provider_list[0] if provider_list else "Unknown provider"

        # Object type (IMAGE, VIDEO, TEXT, SOUND, etc.)
        obj_type = item.get("type") or "Unknown type"

        # A short description
        desc_list = item.get("dcDescription") or []
        desc_str = desc_list[0] if desc_list else "No description available."

        # Year (if present)
        year_list = item.get("year") or []
        year_str = year_list[0] if year_list else "N/A"

        # Construct a child-friendly summary
        # Feel free to reword for an even more "kid-friendly" vibe
        summary_text = (
            f"{idx}) **Title**: {title_str}\n"
            f"   **Where it’s from**: {provider_str}\n"
            f"   **Type of item**: {obj_type}\n"
            f"   **Approx. Year**: {year_str}\n"
            f"   **Fun Fact/Description**: {desc_str}\n"
        )
        result_lines.append(summary_text)

    # Combine the lines into a single string
    intro = "Here are some cool things I found in Europeana:\n"
    return intro + "\n".join(result_lines)

########################################
# EUROPEANA TOOL
########################################
EUROPEANA_API_KEY = "vievinatme"  #

@tool
def query_europeana(query: str) -> str:
    """
    A tool that queries the Europeana Search API for a given query
    and returns up to 5 results in a kid-friendly summary.
    
    Args:
        query: A string representing the search term (e.g. 'Van Gogh')
        
    Returns:
        A string summary describing up to 5 Europeana items in a child-friendly format.
    """
    endpoint = "https://api.europeana.eu/record/v2/search.json"
    params = {
        "query": query,
        "wskey": EUROPEANA_API_KEY,
        "rows": 5,
    }
    try:
        response = requests.get(endpoint, params=params)
        data = response.json()
        
        if response.status_code != 200:
            return f"Oops, something went wrong: {data.get('error', 'Unknown HTTP error')}"

        return restructure_europeana_response(data)

    except Exception as e:
        return f"Error calling Europeana API: {str(e)}"

########################################
# TIME TOOL
########################################
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """
    A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    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)}"


########################################
# WIKIPEDIA CULTURAL INFO
########################################
@tool
def get_cultural_info(topic: str) -> str:
    """
    A tool that retrieves cultural or general info from Wikipedia for a given topic.
    Args:
        topic: A string representing the subject to lookup (e.g., 'Renaissance art').
    Returns:
        A short summary text from Wikipedia for the specified topic.
    """
    try:
        url = (
            "https://en.wikipedia.org/w/api.php?"
            "action=query&prop=extracts&exintro&explaintext&format=json&titles=" + topic
        )
        response = requests.get(url)
        data = response.json()

        pages = data.get("query", {}).get("pages", {})
        if not pages:
            return f"I couldn't find anything on Wikipedia for '{topic}'."

        page_id = next(iter(pages))
        page_content = pages[page_id]
        if "missing" in page_content:
            return f"No page found on Wikipedia for topic: {topic}"

        extract = page_content.get("extract", "")
        if not extract:
            return f"No extract available for '{topic}'."

        return extract.strip()

    except Exception as e:
        return f"Error retrieving cultural info for '{topic}': {str(e)}"


########################################
# FINAL ANSWER TOOL (REQUIRED)
########################################
final_answer = FinalAnswerTool()

########################################
# MODEL 
########################################
model = HfApiModel(
    max_tokens=1024,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',  # or your chosen HF endpoint
    custom_role_conversions=None
)

########################################
# OPTIONAL: IMAGE GENERATION TOOL
########################################
try:
    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
except Exception as e:
    print("Error loading text-to-image tool:", e)
    image_generation_tool = None

########################################
# PROMPT TEMPLATES (for kids style)
########################################
this_dir = os.path.dirname(os.path.abspath(__file__))
prompts_path = os.path.join(this_dir, "prompts.yaml")
with open(prompts_path, 'r', encoding='utf-8') as stream:
    prompt_templates = yaml.safe_load(stream)

########################################
# BUILD THE AGENT
########################################
tools_list = [
    final_answer,  
    DuckDuckGoSearchTool(),
    VisitWebpageTool(),
    get_current_time_in_timezone,
    get_cultural_info,
    query_europeana,  # <--- Our new Europeana tool
]

if image_generation_tool:
    tools_list.append(image_generation_tool)

agent = CodeAgent(
    model=model,
    tools=tools_list,
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="KidsMuseumAgent",
    description="A friendly museum assistant that explains art, history, and culture in kid-friendly language.",
    prompt_templates=prompt_templates
)

if __name__ == "__main__":
    # (OPTIONAL) Launch Gradio Chat UI
    from Gradio_UI import GradioUI
    GradioUI(agent).launch()