from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from tools.visit_webpage import VisitWebpageTool
import os
import IPython.display as display
from typing import Union

import io
import base64
from PIL import Image  # Only needed if 'image' is a PIL Image object

from Gradio_UI import GradioUI


@tool
def generate_surfer_image(next_tide: str) -> str:
    """ A tool that helps to answer users question about whether to go surfing or not. 
        It generates an image of a surfer in the ocean or chilling on Falesia beach and waiting, depending on whether the next tide is high or low.
        If the closest tide is high, generate the surfer in the ocean surfing. If the closest tide is low, generate the surfer chilling. 
        The image should be displayed as an answer to the user's question.
    Args:
        next_tide: A string representing whether next tide low or high.
    
    Returns:
        A string representing the generated image (e.g., a URL or image data), or an error message.
    """
    action = 'chilling on Falesia beach'
    if next_tide == 'high':
        action = 'surfing in Algarve, Falesia'
    prompt = f"Generate a surfer {action}"
    try:
        image_result = image_generation_tool(prompt) 
        # image_byte_arr = io.BytesIO()
        # image_result.save(image_byte_arr, format='JPEG')
        # image_byte_arr = image_byte_arr.getvalue()
        # base64_image = base64.b64encode(image_byte_arr).decode('utf-8')
        
        # return f"<img src='data:image/png;base64,{base64_image}' alt='Generated Image'>" #return html image tag.
        return image_result
        
    except Exception as e:
        return f"Error generating image: {str(e)}"


@tool
def display_image_tool(image: str) -> None:
    """
    Display an image from a URL or a Base64 string.

    Args:
        image: The image to display. It must be: a URL string (starting with "http") or a Base64-encoded image string (starting with "data:image").
    Returns:
        None: The function does not return anything. It displays the image.

    Raises:
        ValueError: If the string format is not a valid image URL or Base64 string.
    """
    if image.startswith("http"):
        # Display image from URL
        display.display(display.HTML(f'<img src="{image}" width="512">'))

    elif image.startswith("data:image"):
        # Display Base64-encoded image
        display.display(display.HTML(f'<img src="{image}" width="512">'))

    else:
        raise ValueError("Invalid image format: Expected an image URL or Base64 string.")





@tool
def tide_time(date: str) -> str:
    """ A tool that checks when the time of the next high tide in Falesia beach in Algarve using DuckDuckGo search.
        Args:
            date: A string representing a valid time in Lisbon time zone.
        Returns: A string representing a valid time of high tide.
    """
    query = f"What is the next high tide time in Falesia beach, Algarve? The current time is {date}"
    try:
        duck_tool = DuckDuckGoSearchTool()
        results = duck_tool(query)
    except Exception as e:
        return f'An error occured while searching: {e}'
        
    if isinstance(results, list):
            if len(results) > 0:
                first_result = results[0]
                if isinstance(first_result, dict):
                    return first_result.get("snippet", "No snippet available for the first result.")
                elif isinstance(first_result, str):
                    return first_result
            return "No tide info found for that date."
    elif isinstance(results, str):
        return results
    else:
        return "No tide info found for that date."
        

@tool
def get_tide_time() -> str:
    """ A tool that visits the given webpage and checks the time of the next tide in Falesia beach in Algarve.
        Returns: A string representing a valid time of the next tide.
    """
    url = 'https://www.tideschart.com/Portugal/Faro/Albufeira/Praia-da-Falesia/'
    try:
        web_tool = VisitWebpageTool()
        response = web_tool(url)
        return response
    except Exception as ex:
        return f'An error occured while retrieving: {e}'
        
    


@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., 'Lisbon').
    """
    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()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

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
model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
custom_role_conversions=None,
)


# Import tool from Hub
# image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
image_generation_tool = load_tool(
    "agents-course/text-to-image",
    model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud",
    trust_remote_code=True,
    token=os.getenv('HF_TOKEN')
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[get_current_time_in_timezone,get_tide_time,generate_surfer_image,display_image_tool,final_answer], ## 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()