import inspect
import asyncio
from typing import get_type_hints

class ToolParameter:
    """Describes a single input parameter for a tool.

    Mirrors the JSON-schema style used by OpenAI / llama-cpp function calling.

    Attributes:
        name: Parameter name.
        type: JSON-schema type string, e.g. "string", "number", "boolean".
        description: Human-readable description of the parameter.
        required: Whether the parameter is required.
        enum: Optional list of allowed values.
    """

    def __init__(self, name, type="string", description="", required=True, enum=None):
        self.name = name
        self.type = type
        self.description = description
        self.required = required
        self.enum = enum

    def to_dict(self):
        d = {
            "type": self.type,
            "description": self.description,
        }
        if self.enum is not None:
            d["enum"] = self.enum
        return d

    def __repr__(self):
        return f"ToolParameter(name={self.name!r}, type={self.type!r})"


class ToolSpec:
    """Declarative specification for a tool callable by an LLM.

    A ToolSpec describes *what* a tool does and *what* parameters it accepts.
    It is used at inference time to build the tool-description payload sent to
    the model so the model can emit structured tool-call requests.

    Inspired by ``langchain.tools.StructuredTool`` and OpenAI's
    ``tools`` / ``functions`` request format.

    Attributes:
        name: Unique identifier for the tool.
        description: Short human-readable description shown to the model.
        parameters: List of ``ToolParameter`` objects.
        # execution
        func: Optional callable that executes the tool.  When present the
              orchestrator can dispatch a ``ToolCall`` automatically.
    """

    def __init__(self, name, description="", parameters=None, func=None):
        self.name = name
        self.description = description
        self.parameters = parameters or []
        self.func = func
    
    def clone (self):
        return ToolSpec(
            name = self.name,
            description = self.description,
            parameters = self.parameters,
            func = self.func,
        )

    # -- serialization helpers -------------------------------------------

    def to_openai_dict(self):
        """Return a dict in OpenAI chat-completions ``tools`` format."""
        properties = {}
        required = []
        for p in self.parameters:
            properties[p.name] = p.to_dict()
            if p.required:
                required.append(p.name)

        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": {
                    "type": "object",
                    "properties": properties,
                    "required": required,
                },
            },
        }

    def to_replicate_dict(self):
        """Return a dict suitable for Replicate model input.

        Replicate models do not have a standard tool-calling protocol, so we
        return a simplified representation that can be injected into the
        prompt as text.
        """
        param_lines = []
        for p in self.parameters:
            req = " (required)" if p.required else " (optional)"
            param_lines.append(f"  - {p.name} ({p.type}){req}: {p.description}")
        params_text = "\n".join(param_lines) if param_lines else "  (none)"
        return {
            "name": self.name,
            "description": self.description,
            "parameters_text": f"{self.name}:\n{params_text}",
        }

    def execute(self, **kwargs):
        """Synchronously run the tool with the given keyword arguments."""
        if self.func is None:
            raise RuntimeError(f"ToolSpec '{self.name}' has no callable (func is None)")
        return self.func(**kwargs)

    async def execute_async(self, **kwargs):
        """Asynchronously run the tool with the given keyword arguments."""
        if self.func is None:
            raise RuntimeError(f"ToolSpec '{self.name}' has no callable (func is None)")
        if asyncio.iscoroutinefunction(self.func):
            return await self.func(**kwargs)
        return self.func(**kwargs)

    def __repr__(self):
        return f"ToolSpec(name={self.name!r}, params={[p.name for p in self.parameters]})"

class Tool:
    """Base class for tools that implement their logic in a 'run' method.

    Subclasses should override 'run' to provide the tool's functionality.
    
    The Tool class provides a method to_spec() that automatically generates
    a ToolSpec from the class name, docstring, and run method signature.
    This allows tools to be used with LLM function calling without manually
    creating ToolSpec objects. The generated ToolSpec is cached at initialization.
    """

    def __init__(self, name=None):
        # Generate and cache the ToolSpec during initialization
        self.name = name or self.__class__.tool_name()
        self._spec = self._generate_spec()

    def run(self, **kwargs):
        """Execute the tool. To be overridden by subclasses."""
        raise NotImplementedError("Tool subclasses must override the 'run' method.")
    
    @classmethod
    def tool_name(cls):
        return ''.join(['_'+c.lower() if c.isupper() else c for c in cls.__name__]).lstrip('_')


    def _generate_spec(self):
        """Generate a ToolSpec from this tool instance.
        
        Automatically creates a ToolSpec by:
        - Using the class name (converted to snake_case) as the tool name
        - Using the class docstring as the description
        - Inspecting the run method signature to create parameters
        """
        
        # Use the class docstring as the description
        description = inspect.getdoc(self.__class__) or ""
        
        # Get parameters from the run method signature
        parameters = []
        run_signature = inspect.signature(self.run)
        type_hints = get_type_hints(self.run)
        
        for param_name, param in run_signature.parameters.items():
            # Skip 'self' parameter
            if param_name == 'self':
                continue
                
            # Get type information
            param_type = type_hints.get(param_name, "string")
            if isinstance(param_type, type):
                type_name = param_type.__name__.lower()
            else:
                type_name = str(param_type).lower()
                
            # Check if parameter is required (no default or default is None)
            is_required = param.default is inspect.Parameter.empty or param.default is None
            
            # Create ToolParameter
            tool_param = ToolParameter(
                name=param_name,
                type=type_name,
                description="",
                required=is_required
            )
            parameters.append(tool_param)
        
        return ToolSpec(
            name=self.name,
            description=description,
            parameters=parameters,
            func=self.run,
        )

    def get_spec(self):
        """Return the cached ToolSpec for this tool."""
        return self._spec

    def __repr__(self):
        return f"Tool(name={self._spec.name!r}, params={[p.name for p in self._spec.parameters]})"

class ToolCall:
    """Represents a single tool invocation request emitted by the model.

    Attributes:
        id: Unique call identifier (used by OpenAI / llama-cpp protocol).
        tool_name: Name of the tool to invoke.
        arguments: Dict of argument name → value.
        result: String content returned by the tool.
        is_error: Whether the result represents an error.
    """

    def __init__(self, tool_name, arguments={}, id=None):
        self.id = id or ""
        self.tool_name = tool_name
        self.arguments = arguments

        self.result = None

    def __repr__(self):
        id_prefix = self.id + ', ' if self.id else ''
        return f"ToolCall({id_prefix}tool={self.tool_name!r}, args={self.arguments!r})"

    def call_message(self):
        raise Exception("TODO: Llama message for tool call")

    def result_message(self):
        from inference import LlamaMessage
        return LlamaMessage("tool", self.result)


# ---------------------------------------------------------------------------
# Helpers to format tool specs per engine
# ---------------------------------------------------------------------------

def format_tools_for_engine(tools, engine):
    """Convert a list of ``ToolSpec`` objects for the given engine.

    Returns a value suitable to pass as the ``tools`` kwarg to the underlying
    chat-completion call.
    """
    specs = [t for t in tools]

    if engine in ("llama-cpp", "ollama"):
        return [s.to_openai_dict() for s in specs]
    elif engine == "replicate":
        return [s.to_replicate_dict() for s in specs]
    else:
        return [s.to_openai_dict() for s in specs]


def build_tool_prompt_prefix(tools):
    """Build a text prefix describing available tools for engines that rely
    on prompt-injection (e.g. Replicate) rather than native tool-calling."""
    if not tools:
        return ""
    parts = ["You have access to the following tools:"]
    for t in tools:
        spec = t
        info = spec.to_replicate_dict()
        parts.append(info["parameters_text"])
    parts.append("")
    return "\n".join(parts)