import base64
import os
from smolagents import Tool
from openai import OpenAI

class AnalyzeImageTool(Tool):
    name = "analyze_image_tool"
    description = """This tool performs a custom analysis of the provided image and returns the corresponding result."""
    inputs = {
        "image_path": {"type": "string", "description": "Image path"},
        "task": {"type": "string", "description": "Task to perform on the image, be detailed and clear"},
    }
    output_type = "string"

    def __init__(self):
        super().__init__()
        self.model_id = "gpt-4.1-mini"

    def forward(self, image_path: str, task: str) -> str:
        """
        Analyze the image at `image_path` according to `task` and return the textual result.
        """
        header = "Image analysis result:\n\n"
        llm_instruction = (
            "You are a highly capable image analysis tool, designed to examine images and deliver detailed descriptions, "
            "insights, and relevant interpretations based on the task at hand.\n\n"
            "Approach the task methodically and provide a thorough and well-reasoned response to the following:\n\n---\nTask:\n"
            f"{task}\n\n"
        )
        try:
            return header + self._analyze_with_openai(image_path, llm_instruction)
        except Exception as e:
            return f"Error analyzing image: {e}."

    def _analyze_with_openai(self, image_path: str, task: str) -> str:
        client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL"))

        with open(image_path, "rb") as f:
            encoded_image = base64.b64encode(f.read()).decode("utf-8")

        payload = [
            {
                "role": "user",
                "content": [
                    {"type": "input_text", "text": task},
                    {"type": "input_image", "image_url": f"data:image/jpeg;base64,{encoded_image}"},
                ],
            }
        ]
        response = client.responses.create(model=self.model_id, input=payload)
        return response.output[0].content[0].text