from typing import Any, Optional
from smolagents.tools import Tool
from PIL import Image
import os
import datetime

class FinalAnswerTool(Tool):
    name = "final_answer"
    description = "Provides a final answer to the given problem. Can include a textual response and/or a path to a generated image."
    
    # CORRECTED: Add "nullable": True for both 'text' and 'image_path'
    inputs = {
        'text': {'type': 'string', 'description': 'The final textual answer to the problem', 'nullable': True},
        'image_path': {'type': 'string', 'description': 'Optional path to a generated image file (e.g., .webp, .png, .jpg) to be displayed.', 'nullable': True}
    }
    
    output_type = "string" # Output will be a Markdown string

    def forward(self, text: Optional[str] = None, image_path: Optional[str] = None) -> str:
        """
        Processes the final answer, potentially embedding an image from a given path.
        """
        output_content = []

        if image_path and os.path.exists(image_path):
            image_display_path = image_path
            if image_path.lower().endswith(".webp"):
                try:
                    img = Image.open(image_path)
                    temp_dir = os.path.dirname(image_path) if os.path.dirname(image_path) else "/tmp"
                    os.makedirs(temp_dir, exist_ok=True)
                    
                    base_name = os.path.basename(image_path)
                    name_without_ext = os.path.splitext(base_name)[0]
                    temp_png_path = os.path.join(
                        temp_dir,
                        f"{name_without_ext}_{datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')}.png"
                    )
                    
                    img.save(temp_png_path, "PNG")
                    image_display_path = temp_png_path
                    print(f"[DEBUG] Converted {image_path} to {temp_png_path} for display.")
                except Exception as e:
                    print(f"[ERROR] Failed to convert image {image_path} to PNG: {e}. Attempting to use original path.")
        
            output_content.append(f"![Generated Image]({image_display_path})")

        if text:
            output_content.append(text)
        
        if not output_content:
            # If both are None, provide a default message
            return "No final answer provided."

        return "\n\n".join(output_content)

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