# from typing import Any, Optional
# from smolagents.tools import 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 forward(self, answer: Any) -> Any:
#         # ✅ If the output is an image, convert it to base64 HTML so Gradio displays it
#         if isinstance(answer, Image.Image):
#             buffer = io.BytesIO()
#             answer.save(buffer, format="PNG")
#             img_str = base64.b64encode(buffer.getvalue()).decode()
#             html = f'<img src="data:image/png;base64,{img_str}" width="512"/>'
#             return html

#         # ✅ If it’s an AgentImage (from smolagents types)
#         try:
#             from smolagents.agent_types import AgentImage
#             if isinstance(answer, AgentImage):
#                 buffer = io.BytesIO()
#                 answer.image.save(buffer, format="PNG")
#                 img_str = base64.b64encode(buffer.getvalue()).decode()
#                 html = f'<img src="data:image/png;base64,{img_str}" width="512"/>'
#                 return html
#         except ImportError:
#             pass

#         # ✅ Otherwise, just return the plain answer
#         return str(answer)

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


from typing import Any
from smolagents.tools import Tool
from PIL import Image  # ✅ this import is crucial
import io
import base64

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

    def forward(self, answer: Any) -> Any:
        # ✅ Case 1: Answer is a plain PIL image
        if isinstance(answer, Image.Image):
            buffer = io.BytesIO()
            answer.save(buffer, format="PNG")
            img_str = base64.b64encode(buffer.getvalue()).decode()
            html = f'<img src="data:image/png;base64,{img_str}" width="512"/>'
            return html

        # ✅ Case 2: Answer is a smolagents.AgentImage (sometimes wrapped)
        try:
            from smolagents.agent_types import AgentImage
            if isinstance(answer, AgentImage):
                buffer = io.BytesIO()
                answer.image.save(buffer, format="PNG")
                img_str = base64.b64encode(buffer.getvalue()).decode()
                html = f'<img src="data:image/png;base64,{img_str}" width="512"/>'
                return html
        except ImportError:
            pass

        # ✅ Default fallback: return text as-is
        return str(answer)

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