from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool, AgentText
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
import os
import wikipediaapi
import json
import re

from Gradio_UI import GradioUI


@tool
def get_country_info(country_name: str) -> AgentText:
    """Fetches detailed information about a country using the REST Countries API and Wikipedia.
    Args:
        country_name: The name of the country to look up (e.g., 'France').
    Returns:
        AgentText: Formatted country information with embedded images
    """
    try:
        # Get the API key from environment variable
        GEOAPIFY_API_KEY = os.getenv('GEOAPIFY_API_KEY')

        if not GEOAPIFY_API_KEY:
            raise ValueError("GEOAPIFY_API_KEY environment variable is not set")

        # Make the API request to REST Countries
        url = f"https://restcountries.com/v3.1/name/{country_name}?fullText=true"
        response = requests.get(url)
        response.raise_for_status()
        
        # Parse the response
        data = response.json()[0]
        
        # Safely extract nested information with fallbacks
        def get_nested(obj, *keys, default="N/A"):
            try:
                for key in keys:
                    obj = obj[key]
                return obj
            except (KeyError, TypeError, IndexError):
                return default

        # Extract key information with safe fallbacks
        info = {
            "official_name": get_nested(data, "name", "official"),
            "capital": get_nested(data, "capital", 0),
            "region": get_nested(data, "region"),
            "subregion": get_nested(data, "subregion"),
            "population": f"{get_nested(data, 'population', default=0):,}",
            "languages": ", ".join(get_nested(data, "languages", default={}).values()),
            "flag": get_nested(data, "flag"),
            "flag_url": get_nested(data, "flags", "png"),
            "latlng": get_nested(data, "latlng", default=[]),
            "borders": get_nested(data, "borders", default=[])
        }

        # Safely extract currencies
        currencies = get_nested(data, "currencies", default={})
        currency_list = []
        for currency_code, currency_data in currencies.items():
            name = get_nested(currency_data, "name", default="")
            symbol = get_nested(currency_data, "symbol", default="")
            if name and symbol:
                currency_list.append(f"{name} ({symbol})")
            elif name:
                currency_list.append(name)
        info["currencies"] = ", ".join(currency_list) if currency_list else "N/A"

        # Get full names of bordering countries
        border_countries = []
        if info["borders"]:
            for border_code in info["borders"]:
                try:
                    border_url = f"https://restcountries.com/v3.1/alpha/{border_code}"
                    border_response = requests.get(border_url)
                    border_response.raise_for_status()
                    border_data = border_response.json()[0]
                    border_countries.append(get_nested(border_data, "name", "common"))
                except:
                    border_countries.append(border_code)
        
        borders_text = "None" if not border_countries else ", ".join(sorted(border_countries))

        # Get Wikipedia information
        wiki = wikipediaapi.Wikipedia('CountryInfoBot (contact@example.com)', 'en')
        wiki_page = wiki.page(country_name)
        
        wiki_info = {'intro': 'No additional information available.'}
        if wiki_page.exists():
            wiki_info['intro'] = wiki_page.summary.split('\n')[0]

        # Create static map URL using coordinates
        if info["latlng"] and len(info["latlng"]) == 2:
            lat, lng = info["latlng"]
            static_map_url = (
                f"https://maps.geoapify.com/v1/staticmap"
                f"?style=osm-bright"
                f"&width=600"
                f"&height=400"
                f"&center=lonlat:{lng},{lat}"
                f"&zoom=4"
                f"&apiKey={GEOAPIFY_API_KEY}"
            )
        else:
            static_map_url = None

        # Format the response as a specially marked section for the agent
        response_text = f"""<country_info_start>
# {country_name} {info['flag']}

| Flag | Map |
|:---:|:---:|
| ![Flag of {country_name}]({info['flag_url']}) | ![Map of {country_name}]({static_map_url if static_map_url else 'Not available'}) |

## Basic Information
- Official Name: {info['official_name']}
- Capital: {info['capital']}
- Region: {info['region']} ({info['subregion']})
- Population: {info['population']}
- Languages: {info['languages']}
- Currencies: {info['currencies']}
- Bordering Countries: {borders_text}

## Overview
{wiki_info['intro']}
<country_info_end>"""

        return AgentText(response_text)

    except requests.exceptions.RequestException as e:
        return AgentText(f"Error fetching country information: {str(e)}")
    except Exception as e:
        return AgentText(f"Unexpected error: {str(e)}")


# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_custom_tool(arg1:str, arg2:int)-> str: #it's important to specify the return type
    # Keep this format for the description / args / args description but feel free to modify the tool
    """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:
    """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:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that 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)}"


final_answer = FinalAnswerTool()
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct', # it is possible that this model may be overloaded
custom_role_conversions=None,
token=os.getenv('hf_token')
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        get_country_info], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()