import asyncio
import json
import logging
from typing import Optional

from openai import AsyncOpenAI

from ..config import AgentConfig
from .manager import ToolManager
from ...components import MessageStorageBase
from ...schemas import Message, ToolCall, RoleType, MessageType
from ...utils.image_utils import is_image_output, extract_source


logger = logging.getLogger(__name__)


class ToolExecutor:
    """Executes tool calls, handles image detection/captioning, and manages concurrent execution."""

    def __init__(
        self,
        tool_manager: ToolManager,
        message_storage: MessageStorageBase,
        client: AsyncOpenAI,
    ):
        self.tool_manager = tool_manager
        self.message_storage = message_storage
        self.client = client

    async def handle_tool_calls(
        self,
        tool_calls: list,
        user_id: str,
        session_id: str,
        input_messages: list,
        max_concurrent_tools: int = AgentConfig.DEFAULT_MAX_CONCURRENT_TOOLS,
    ) -> Optional[tuple[str, str]]:
        """
        Handle tool calls by executing them concurrently with concurrency limit.

        Returns:
            None if no image output, or (image_data, description) tuple.
        """
        if not tool_calls:
            return None

        function_calls = [tc for tc in tool_calls if getattr(tc, "type", None) == "function_call"]
        if not function_calls:
            return None

        semaphore = asyncio.Semaphore(max_concurrent_tools)

        async def execute_with_semaphore(tool_call):
            async with semaphore:
                return await self.execute_single(tool_call, user_id, session_id)

        tasks = [execute_with_semaphore(tc) for tc in function_calls]
        results = await asyncio.gather(*tasks)

        pending_images = []
        pending_descriptions = []

        for tool_messages, image_data, description in results:
            if tool_messages:
                input_messages.extend([msg.to_dict() for msg in tool_messages])
            if image_data:
                pending_images.append(image_data)
            if description:
                pending_descriptions.append(description)

        if pending_images:
            return (
                "\n\n".join(pending_images),
                "\n\n".join(pending_descriptions),
            )

        return None

    async def execute_single(
        self,
        tool_call,
        user_id: str,
        session_id: str,
    ) -> tuple[Optional[list], Optional[str], Optional[str]]:
        """Execute a single tool call and return (messages, image_data, description)."""
        name = getattr(tool_call, "name", None)
        try:
            args = json.loads(getattr(tool_call, "arguments", "{}"))
        except Exception as e:
            logger.error("Tool args parse error: %s", e)
            return None, None, None

        func = self.tool_manager.get_tool(name)
        if not func:
            return None, None, None

        logger.info("Calling tool: %s with args: %s", name, args)

        try:
            result = await func(**args)
        except Exception as e:
            logger.error("Tool call error: %s", e)
            result = f"Tool error: {e}"

        result_str = json.dumps(result, ensure_ascii=False) if isinstance(result, (dict, list)) else str(result)

        # Detect image results
        image_data = None
        model_output = result_str
        if is_image_output(result_str):
            image_data = result_str
            image_src = extract_source(result_str)
            prompt_hint = args.get("prompt", "")
            caption = await self._caption_image(image_src, prompt_hint)
            model_output = (
                f"[Image generated by tool `{name}` and displayed to user. "
                f"Image description: {caption}]"
            )

        tool_call_msg = Message(
            type=MessageType.FUNCTION_CALL,
            role=RoleType.TOOL,
            content=f"Calling tool: `{name}` with args: {args}",
            tool_call=ToolCall(
                call_id=getattr(tool_call, "call_id", ""),
                name=name,
                arguments=json.dumps(args, ensure_ascii=False)
            )
        )

        def _format_preview(text: str) -> str:
            if len(text) <= AgentConfig.TOOL_RESULT_PREVIEW_LENGTH:
                return text
            return text[:AgentConfig.TOOL_RESULT_PREVIEW_LENGTH] + '...'

        tool_res_msg = Message(
            type=MessageType.FUNCTION_CALL_OUTPUT,
            role=RoleType.TOOL,
            content=f"Tool `{name}` result: {_format_preview(result_str)}",
            tool_call=ToolCall(
                call_id=getattr(tool_call, "call_id", "001"),
                output=model_output
            )
        )
        await self.message_storage.add_messages(user_id, session_id, [tool_call_msg, tool_res_msg])

        return [tool_call_msg, tool_res_msg], image_data, model_output if image_data else None

    async def _caption_image(self, image_data_uri: str, prompt_hint: str = "") -> str:
        """Use a vision model to generate a detailed description of an image."""
        caption_prompt = AgentConfig.IMAGE_CAPTION_PROMPT
        if prompt_hint:
            caption_prompt += f"\n\nOriginal generation prompt: \"{prompt_hint}\""

        try:
            response = await self.client.responses.create(
                model=AgentConfig.IMAGE_CAPTION_MODEL,
                input=[{
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": caption_prompt},
                        {"type": "input_image", "image_url": image_data_uri},
                    ]
                }],
            )
            caption = getattr(response, "output_text", "") or ""
            if caption.strip():
                return caption.strip()
        except Exception as e:
            logger.warning("Image captioning failed, falling back to prompt-based description: %s", e)

        if prompt_hint:
            return f"Generated image based on prompt: \"{prompt_hint}\""
        return "An image was generated and displayed to the user."
