from typing import Any, Optional
from smolagents.tools import Tool
import io
import openpyxl
import os
from smolagents import tool
import requests
from PIL import Image


simplify_system_message = """
You are a general AI assistant. I will give you a question and answer, you should simplify answer. Your answer should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
"""     

@tool
def ToolReverseString(text: str) -> str:
    """
    Use this tool only if you encounter text that seems to be written backwards
    Args:
        text: The string to reverse.
    """
    return text[::-1]


@tool
def ToolReadFiles(filepath: str) -> str:
    """
    Downloads a .py or .xlsx file from a remote URL and returns its contents as plain text.
    Raises a recoverable exception if the file does not end with .py or .xlsx.
    Args:
        filepath: The path to the Python (.py) or Excel (.xlsx) file.
    """
    root_url = "https://agents-course-unit4-scoring.hf.space/files/"
    # Strip the file extension from the url before downloading
    import os
    base, ext = os.path.splitext(filepath)
    url = root_url + base

    if filepath.endswith('.py'):
        response = requests.get(url)
        if response.status_code != 200:
            raise Exception(f"Recoverable: Failed to download file from {url}")
        return response.text

    elif filepath.endswith('.xlsx'):
        response = requests.get(url)
        if response.status_code != 200:
            raise Exception(f"Recoverable: Failed to download file from {url}")

        wb = openpyxl.load_workbook(io.BytesIO(response.content), data_only=True)
        result = []
        for sheet in wb.worksheets:
            result.append(f"# Sheet: {sheet.title}")
            for row in sheet.iter_rows(values_only=True):
                result.append(','.join([str(cell) if cell is not None else '' for cell in row]))
        return '\n'.join(result)

    else:
        raise Exception("Recoverable: Only .py and .xlsx files can be read with this tool.")


@tool
def ToolDownloadImage(filepath: str) -> str:
    """
    Downloads an image file (.png, .jpg, .jpeg) from a remote URL and returns useful information about the image.
    This includes the image URL and basic metadata like dimensions and format.
    Raises a recoverable exception if the file is not a supported image type.
    Args:
        filepath: The path to the image file.
    """
    root_url = "https://agents-course-unit4-scoring.hf.space/files/"
    base, ext = os.path.splitext(filepath)
    url = root_url + base
    
    if ext.lower() in ['.png', '.jpg', '.jpeg']:
        response = requests.get(url)
        if response.status_code != 200:
            raise Exception(f"Recoverable: Failed to download image from {url}")
        
        # Get image metadata using Pillow
        try:
            
            img = Image.open(io.BytesIO(response.content))
            width, height = img.size
            format = img.format
            mode = img.mode
            
            # Return useful information about the image
            return f"Image URL: {url}\nFormat: {format}\nDimensions: {width}x{height}\nMode: {mode}"
        except ImportError:
            # Fallback if PIL is not available
            content_type = response.headers.get('Content-Type', 'unknown')
            content_length = response.headers.get('Content-Length', 'unknown')
            return f"Content-Type: {content_type}\nSize: {content_length} bytes"
    else:
        raise Exception("Recoverable: Only .png, .jpg, and .jpeg files can be processed with this tool.")

#class FinalAnswerTool(Tool):
#    name = "final_answer"
#    description = "Provides a final answer to the given problem."
#    inputs = {'answer': {'type': 'any', 'description': 'The final answer to the problem'}}
#    output_type = "any"

#    def forward(self, answer: Any) -> Any:
#        return answer

#    def __init__(self, *args, **kwargs):
#        self.is_initialized = False


class FinalAnswerTool(Tool):
    name = "simplify_answer"
    description = (
        "Simplify the final answer."
        "Be sure to use this as the last step to prepare the final answer."
    )
    inputs = {
        "question": {
            "type": "string",
            "description": "The question from the user.",
        },
        "answer": {
            "type": "string",
            "description": "The existing answer to be simplified.",
        }        
    }
    output_type = "string"

    def __init__(self, api_model):
        super().__init__()
        self.api_model = api_model

    def forward(self, question: str, answer: str) -> str:
        try:
            response_message = self.api_model(messages=[
                {"role": MessageRole.SYSTEM, "content": [{"type": "text", "text": simplify_system_message }]},
                {"role": MessageRole.ASSISTANT, "content": [{"type": "text", "text": f"QUESTION: {question}"}]},
                {"role": MessageRole.ASSISTANT, "content": [{"type": "text", "text": f"ANSWER: {answer}"}]}
            ])
            return response_message.content

        except requests.exceptions.Timeout:
            return "The request timed out. Please try again later or check the URL."
        except RequestException as e:
            return f"Error fetching the webpage: {str(e)}"
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"