"""ToolService — unified tool management for Fluxibly.

Central service for registering, translating, and executing tools
across all LLM providers. Handles:

- Function tools (Python callables)
- MCP tools (via MCPClientManager)
- Provider-specific tools (web search, file search, shell, computer use)
- Cross-provider tool format translation
"""

from __future__ import annotations

import asyncio
import json
from collections.abc import Callable
from typing import Any

from fluxibly.logging import Logger
from fluxibly.schema.tools import ToolCall, ToolResult
from fluxibly.tools.mcp_manager import MCPClientManager


class ToolService:
    """Central service for managing and executing all tool types.

    Responsibilities:
    1. Register tools of all types (function, MCP, web_search, etc.)
    2. Translate unified tool format -> provider-specific format
    3. Execute tool calls and return ToolResult
    4. Manage MCP server lifecycles
    """

    def __init__(self) -> None:
        self.mcp_manager: MCPClientManager | None = None
        self.function_registry: dict[str, Callable[..., Any]] = {}
        self.tools: list[dict[str, Any]] = []
        self.skills: list[dict[str, Any]] = []
        self.root_tool_path: str | None = None
        self.logger = Logger(component="tools")

    # ══════════════════════════════════════════════════════════════
    # Tool Registration
    # ══════════════════════════════════════════════════════════════

    def register_function(
        self,
        name: str,
        fn: Callable[..., Any],
        description: str,
        parameters: dict[str, Any],
    ) -> None:
        """Register a Python function as a tool.

        Args:
            name: Tool name.
            fn: Callable (sync or async).
            description: Tool description for LLM.
            parameters: JSON Schema for function arguments.
        """
        self.function_registry[name] = fn
        self.tools.append(
            {
                "type": "function",
                "function": {
                    "name": name,
                    "description": description,
                    "parameters": parameters,
                },
            }
        )

    def register_function_from_schema(
        self,
        schema: dict[str, Any],
        fn: Callable[..., Any],
    ) -> None:
        """Register a function tool from a pre-built OpenAI-format schema.

        Idempotent: skips registration if a tool with the same name is
        already registered.  This allows ``build_tools()`` to call this on
        every ``forward()`` without duplicating entries.

        Args:
            schema: OpenAI tool dict (``{"type": "function", "function": ...}``).
            fn: Callable (sync or async).
        """
        name = schema["function"]["name"]
        if name not in self.function_registry:
            self.function_registry[name] = fn
            self.tools.append(schema)

    async def register_mcp_server(
        self, server_name: str, config: dict[str, Any]
    ) -> None:
        """Register and connect to an MCP server.

        Discovers tools from the server and adds them as function tools.
        """
        if self.mcp_manager is None:
            self.mcp_manager = MCPClientManager()

        await self.mcp_manager.connect_server(server_name, config)

        # Add discovered MCP tools as function tools (all LLMs understand functions)
        for tool in self.mcp_manager.get_server_tools(server_name):
            self.tools.append(
                {
                    "type": "function",
                    "function": {
                        "name": tool["name"],
                        "description": tool["description"],
                        "parameters": tool["schema"],
                    },
                    "_source": "mcp",
                    "_server": server_name,
                }
            )

    def register_web_search(
        self, config: dict[str, Any] | None = None
    ) -> None:
        """Register web search capability."""
        self.tools.append(
            {
                "type": "web_search",
                "web_search": {"enabled": True, **(config or {})},
            }
        )

    def register_file_search(
        self, vector_store_ids: list[str], config: dict[str, Any] | None = None
    ) -> None:
        """Register file search capability."""
        self.tools.append(
            {
                "type": "file_search",
                "file_search": {
                    "vector_store_ids": vector_store_ids,
                    **(config or {}),
                },
            }
        )

    def register_shell(self, config: dict[str, Any] | None = None) -> None:
        """Register shell/code execution capability."""
        self.tools.append(
            {
                "type": "shell",
                "shell": {"enabled": True, **(config or {})},
            }
        )

    def register_computer_use(
        self,
        display_width: int = 1024,
        display_height: int = 768,
        config: dict[str, Any] | None = None,
    ) -> None:
        """Register computer use capability."""
        self.tools.append(
            {
                "type": "computer_use",
                "computer_use": {
                    "display_width": display_width,
                    "display_height": display_height,
                    **(config or {}),
                },
            }
        )

    def register_sandboxed_shell(
        self,
        session: Any,
        tool_name: str = "shell",
    ) -> None:
        """Register a sandboxed shell execution tool.

        Creates a function tool that wraps ``SandboxSession.run()``.
        The tool appears as a regular function tool to the LLM.

        Args:
            session: A ``SandboxSession`` instance from ``fluxibly.tools.sandbox``.
            tool_name: Name for the tool (default: "shell").
        """
        from fluxibly.tools.sandbox import create_shell_tool

        schema, handler = create_shell_tool(session)
        if tool_name != "shell":
            schema["function"]["name"] = tool_name
        self.register_function_from_schema(schema, handler)

    # ══════════════════════════════════════════════════════════════
    # Skill Registration (Agent Skills standard — agentskills.io)
    # ══════════════════════════════════════════════════════════════

    def register_skill(
        self,
        skill_id: str,
        source: str = "custom",
        version: str = "latest",
        **kwargs: Any,
    ) -> None:
        """Register an Agent Skill (agentskills.io standard).

        Idempotent: skips registration if a skill with the same ``skill_id``
        is already registered.

        Skills are versioned bundles of instructions, scripts, and resources.
        Supported natively by OpenAI (Responses API shell) and Anthropic (container).

        Args:
            skill_id: Skill identifier.
                - Anthropic pre-built: "pptx", "xlsx", "docx", "pdf"
                - Custom: "skill_01AbCd..." (Anthropic) or uploaded ID (OpenAI)
            source: Skill source:
                - "anthropic": Pre-built Anthropic skills
                - "openai": OpenAI hosted skills
                - "custom": User-uploaded custom skills
                - "inline": Inline skill (data provided directly)
            version: Version string or "latest".
            **kwargs: Additional skill config (name, description, data, media_type for inline).
        """
        # Idempotent — skip if already registered
        if any(s["skill_id"] == skill_id for s in self.skills):
            return

        skill: dict[str, Any] = {
            "skill_id": skill_id,
            "source": source,
            "version": version,
        }
        skill.update(kwargs)
        self.skills.append(skill)

    def register_inline_skill(
        self,
        name: str,
        description: str,
        data: str,
        media_type: str = "application/zip",
    ) -> None:
        """Register an inline skill with embedded data.

        Args:
            name: Skill name.
            description: Skill description.
            data: Base64-encoded skill archive.
            media_type: MIME type of the data.
        """
        self.skills.append(
            {
                "skill_id": name,
                "source": "inline",
                "version": "inline",
                "name": name,
                "description": description,
                "data": data,
                "media_type": media_type,
            }
        )

    # ══════════════════════════════════════════════════════════════
    # Tool Format Translation
    # ══════════════════════════════════════════════════════════════

    def get_tools_for_provider(
        self, provider: str, api_type: str = "responses"
    ) -> list[dict[str, Any]]:
        """Translate unified tools to provider-specific format.

        Args:
            provider: "openai" | "anthropic" | "gemini" | "langchain" | "litellm"
            api_type: "responses" | "chat" (OpenAI only)

        Returns:
            Provider-formatted tool list ready for LLM.set_tools().
        """
        if provider == "openai":
            if api_type == "responses":
                return self._format_openai_responses(self.tools)
            return self._format_openai_chat(self.tools)
        elif provider == "anthropic":
            return self._format_anthropic(self.tools)
        elif provider == "gemini":
            return self._format_gemini(self.tools)
        else:
            # langchain, litellm, unknown — use OpenAI Chat format
            return self._format_openai_chat(self.tools)

    def _format_openai_responses(
        self, tools: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Format for OpenAI Responses API."""
        formatted: list[dict[str, Any]] = []
        for tool in tools:
            ttype = tool["type"]
            if ttype == "function":
                func = tool["function"]
                formatted.append(
                    {
                        "type": "function",
                        "name": func["name"],
                        "description": func["description"],
                        "parameters": func["parameters"],
                    }
                )
            elif ttype == "web_search":
                entry: dict[str, Any] = {"type": "web_search_preview"}
                ws = tool["web_search"]
                if "search_context_size" in ws:
                    entry["search_context_size"] = ws["search_context_size"]
                if "user_location" in ws:
                    entry["user_location"] = ws["user_location"]
                formatted.append(entry)
            elif ttype == "file_search":
                formatted.append(
                    {
                        "type": "file_search",
                        **{k: v for k, v in tool["file_search"].items()},
                    }
                )
            elif ttype == "shell":
                formatted.append(
                    {
                        "type": "code_interpreter",
                        "container": {"type": "auto"},
                    }
                )
            elif ttype == "computer_use":
                cu = tool["computer_use"]
                formatted.append(
                    {
                        "type": "computer_use_preview",
                        "display_width": cu["display_width"],
                        "display_height": cu["display_height"],
                        "environment": cu.get("environment", "browser"),
                    }
                )
            elif ttype == "mcp":
                mcp = tool["mcp"]
                formatted.append(
                    {
                        "type": "mcp",
                        "server_label": mcp["server_label"],
                        "server_url": mcp.get("server_url", ""),
                        "require_approval": mcp.get(
                            "require_approval", "never"
                        ),
                    }
                )
        return formatted

    def _format_openai_chat(
        self, tools: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Format for OpenAI Chat Completions API (also used by LangChain/LiteLLM)."""
        formatted: list[dict[str, Any]] = []
        for tool in tools:
            if tool["type"] == "function":
                formatted.append(
                    {
                        "type": "function",
                        "function": {
                            "name": tool["function"]["name"],
                            "description": tool["function"]["description"],
                            "parameters": tool["function"]["parameters"],
                        },
                    }
                )
            # Chat API only supports function tools — skip others
        return formatted

    def _format_anthropic(
        self, tools: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Format for Anthropic Messages API."""
        formatted: list[dict[str, Any]] = []
        for tool in tools:
            ttype = tool["type"]
            if ttype == "function":
                func = tool["function"]
                formatted.append(
                    {
                        "name": func["name"],
                        "description": func["description"],
                        "input_schema": func["parameters"],
                    }
                )
            elif ttype == "shell":
                formatted.append(
                    {
                        "name": "bash",
                        "type": "bash_20250124",
                    }
                )
            elif ttype == "computer_use":
                cu = tool["computer_use"]
                formatted.append(
                    {
                        "name": "computer",
                        "type": "computer_20250124",
                        "display_width_px": cu["display_width"],
                        "display_height_px": cu["display_height"],
                    }
                )
            # web_search, file_search: not natively supported by Anthropic
        return formatted

    def _format_gemini(
        self, tools: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Format for Gemini API."""
        function_declarations: list[dict[str, Any]] = []
        other_tools: list[dict[str, Any]] = []

        for tool in tools:
            ttype = tool["type"]
            if ttype == "function":
                func = tool["function"]
                function_declarations.append(
                    {
                        "name": func["name"],
                        "description": func["description"],
                        "parameters": self._convert_json_schema_to_gemini(
                            func["parameters"]
                        ),
                    }
                )
            elif ttype == "web_search":
                other_tools.append({"google_search": {}})
            elif ttype == "shell":
                other_tools.append({"code_execution": {}})
            # computer_use, file_search, mcp: not supported in Gemini

        formatted: list[dict[str, Any]] = []
        if function_declarations:
            formatted.append({"function_declarations": function_declarations})
        formatted.extend(other_tools)
        return formatted

    def _convert_json_schema_to_gemini(
        self, schema: dict[str, Any]
    ) -> dict[str, Any]:
        """Convert JSON Schema types to Gemini format (uppercase types)."""
        type_map = {
            "string": "STRING",
            "number": "NUMBER",
            "integer": "INTEGER",
            "boolean": "BOOLEAN",
            "array": "ARRAY",
            "object": "OBJECT",
        }
        converted: dict[str, Any] = {}
        for key, value in schema.items():
            if key == "type" and isinstance(value, str):
                converted[key] = type_map.get(value, value.upper())
            elif key == "properties" and isinstance(value, dict):
                converted[key] = {
                    k: self._convert_json_schema_to_gemini(v)
                    for k, v in value.items()
                }
            elif key == "items" and isinstance(value, dict):
                converted[key] = self._convert_json_schema_to_gemini(value)
            else:
                converted[key] = value
        return converted

    # ══════════════════════════════════════════════════════════════
    # Tool Execution
    # ══════════════════════════════════════════════════════════════

    async def execute_tool_call(self, tool_call: ToolCall) -> ToolResult:
        """Execute a single tool call and return the result.

        Dispatches to the appropriate executor: registered function or MCP.
        """
        tool_name = tool_call.name
        args = tool_call.arguments
        if isinstance(args, str):
            try:
                args = json.loads(args)
            except json.JSONDecodeError:
                return ToolResult(
                    tool_call_id=tool_call.id,
                    content=f"Invalid JSON arguments for tool '{tool_name}': {args}",
                    is_error=True,
                )

        # 1. Check registered Python functions
        if tool_name in self.function_registry:
            return await self._execute_function(tool_name, args, tool_call.id)

        # 2. Check MCP tools
        if self.mcp_manager and tool_name in self.mcp_manager.tools:
            return await self._execute_mcp(tool_name, args, tool_call.id)

        # 3. Unknown tool
        return ToolResult(
            tool_call_id=tool_call.id,
            content=f"Tool '{tool_name}' not found",
            is_error=True,
        )

    async def execute_tool_calls(
        self, tool_calls: list[ToolCall], max_workers: int = 4
    ) -> list[ToolResult]:
        """Execute a batch of tool calls with concurrency control.

        Args:
            tool_calls: List of tool calls to execute.
            max_workers: Maximum concurrent executions.

        Returns:
            List of ToolResult in same order as input.
        """
        if max_workers <= 1 or len(tool_calls) == 1:
            results: list[ToolResult] = []
            for tc in tool_calls:
                results.append(await self.execute_tool_call(tc))
            return results

        semaphore = asyncio.Semaphore(max_workers)

        async def _run(tc: ToolCall) -> ToolResult:
            async with semaphore:
                return await self.execute_tool_call(tc)

        raw_results: list[ToolResult | BaseException] = await asyncio.gather(
            *[_run(tc) for tc in tool_calls],
            return_exceptions=True,
        )
        final: list[ToolResult] = []
        for i, r in enumerate(raw_results):
            if isinstance(r, BaseException):
                self.logger.error(f"Tool '{tool_calls[i].name}' failed: {r}")
                final.append(
                    ToolResult(
                        tool_call_id=tool_calls[i].id,
                        content=f"Tool execution failed: {r}",
                        is_error=True,
                    )
                )
            else:
                final.append(r)
        return final

    async def _execute_function(
        self, name: str, args: dict[str, Any], call_id: str
    ) -> ToolResult:
        """Execute a registered Python function."""
        fn = self.function_registry[name]
        try:
            if asyncio.iscoroutinefunction(fn):
                result = await fn(**args)
            else:
                result = fn(**args)

            content = (
                json.dumps(result) if not isinstance(result, str) else result
            )
            return ToolResult(tool_call_id=call_id, content=content)

        except Exception as e:
            self.logger.error(f"Function tool '{name}' failed: {e}")
            return ToolResult(
                tool_call_id=call_id,
                content=f"Error executing {name}: {e}",
                is_error=True,
            )

    async def _execute_mcp(
        self, name: str, args: dict[str, Any], call_id: str
    ) -> ToolResult:
        """Execute an MCP tool."""
        try:
            result = await self.mcp_manager.invoke_tool(name, args)  # type: ignore[union-attr]
            return ToolResult(tool_call_id=call_id, content=str(result))
        except Exception as e:
            self.logger.error(f"MCP tool '{name}' failed: {e}")
            return ToolResult(
                tool_call_id=call_id,
                content=f"MCP error for {name}: {e}",
                is_error=True,
            )

    # ══════════════════════════════════════════════════════════════
    # Lifecycle
    # ══════════════════════════════════════════════════════════════

    async def initialize(self) -> None:
        """Initialize all tool backends (MCP servers, etc.)."""
        if self.mcp_manager:
            await self.mcp_manager.initialize()

    async def shutdown(self) -> None:
        """Shutdown all tool backends."""
        if self.mcp_manager:
            await self.mcp_manager.shutdown()

    # ══════════════════════════════════════════════════════════════
    # Accessors
    # ══════════════════════════════════════════════════════════════

    def get_all_tools(self) -> list[dict[str, Any]]:
        """Get all registered tools in unified format."""
        return self.tools

    def get_all_skills(self) -> list[dict[str, Any]]:
        """Get all registered skills in unified format."""
        return self.skills

    def has_tools(self) -> bool:
        """Check if any tools are registered."""
        return bool(self.tools)

    def has_skills(self) -> bool:
        """Check if any skills are registered."""
        return bool(self.skills)

    # ══════════════════════════════════════════════════════════════
    # Skill Format Translation
    # ══════════════════════════════════════════════════════════════

    def get_skills_for_provider(self, provider: str) -> list[dict[str, Any]]:
        """Translate unified skills to provider-specific format.

        Skills are NOT tools — they go to different places in the API:
        - OpenAI: shell tool's environment.skills
        - Anthropic: top-level container.skills parameter
        - Gemini/LangChain/LiteLLM: not supported (returns empty)

        Args:
            provider: "openai" | "anthropic" | "gemini" | "langchain" | "litellm"

        Returns:
            Provider-formatted skill list.
        """
        if provider == "openai":
            return self._format_skills_openai(self.skills)
        elif provider == "anthropic":
            return self._format_skills_anthropic(self.skills)
        return []

    def _format_skills_openai(
        self, skills: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Format skills for OpenAI Responses API (shell environment)."""
        formatted: list[dict[str, Any]] = []
        for skill in skills:
            source = skill.get("source", "custom")

            if source == "inline":
                formatted.append(
                    {
                        "type": "inline",
                        "name": skill.get("name", skill["skill_id"]),
                        "description": skill.get("description", ""),
                        "source": {
                            "type": "base64",
                            "media_type": skill.get(
                                "media_type", "application/zip"
                            ),
                            "data": skill.get("data", ""),
                        },
                    }
                )
            else:
                entry: dict[str, Any] = {
                    "type": "skill_reference",
                    "skill_id": skill["skill_id"],
                }
                version = skill.get("version", "latest")
                if version != "latest":
                    entry["version"] = (
                        int(version) if version.isdigit() else version
                    )
                else:
                    entry["version"] = "latest"
                formatted.append(entry)

        return formatted

    def _format_skills_anthropic(
        self, skills: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Format skills for Anthropic Messages API (container.skills)."""
        formatted: list[dict[str, Any]] = []
        for skill in skills:
            source = skill.get("source", "custom")

            # Map unified source to Anthropic type
            if source == "anthropic":
                skill_type = "anthropic"
            else:
                skill_type = "custom"

            entry: dict[str, Any] = {
                "type": skill_type,
                "skill_id": skill["skill_id"],
            }
            version = skill.get("version")
            if version:
                entry["version"] = version

            formatted.append(entry)

        return formatted
