from __future__ import annotations

import asyncio
import inspect
from types import NoneType
from typing import Any, Union, get_args, get_origin, get_type_hints

from pydantic import BaseModel, ConfigDict, Field

from nexus.protocols.providers.base import ToolDefinition
from nexus.resilience.errors import ToolFailure


class ToolEntry(BaseModel):
    """Canonical tool registration entry."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    name: str
    description: str
    fn: Any = Field(exclude=True)
    schema: dict[str, Any]
    is_async: bool
    source: str = "native"
    tags: list[str] = Field(default_factory=list)


class ToolRegistry:
    """Registry for function tools with schema translation helpers."""

    def __init__(self) -> None:
        self._tools: dict[str, ToolEntry] = {}

    def register(
        self,
        fn,
        name: str | None = None,
        description: str | None = None,
        tags: list[str] | None = None,
    ) -> ToolEntry:
        """Register a callable and infer JSON schema from type hints."""
        tool_name = name or fn.__name__
        desc = description or (inspect.getdoc(fn) or f"Tool {tool_name}")
        schema = self._schema_from_callable(fn)
        entry = ToolEntry(
            name=tool_name,
            description=desc,
            fn=fn,
            schema=schema,
            is_async=inspect.iscoroutinefunction(fn),
            tags=list(tags or []),
        )
        self._tools[tool_name] = entry
        return entry

    def get(self, name: str) -> ToolEntry | None:
        """Fetch tool by name."""
        return self._tools.get(name)

    def list_tools(self, tags: list[str] | None = None) -> list[ToolEntry]:
        """List tools optionally filtered by tags intersection."""
        if not tags:
            return list(self._tools.values())
        required = set(tags)
        return [entry for entry in self._tools.values() if required.intersection(entry.tags)]

    async def call(self, name: str, arguments: dict) -> Any:
        """Invoke a registered tool by name with arguments."""
        entry = self.get(name)
        if entry is None:
            raise ToolFailure(f"Unknown tool '{name}'", tool_name=name, args=arguments)
        try:
            if entry.is_async:
                return await entry.fn(**arguments)
            return await asyncio.to_thread(entry.fn, **arguments)
        except Exception as exc:
            raise ToolFailure(f"Tool '{name}' failed: {exc}", tool_name=name, args=arguments) from exc

    def to_openai_format(self, names: list[str] | None = None) -> list[dict]:
        """Translate registry entries to OpenAI tool format."""
        selected = self._select(names)
        return [
            {
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.schema,
                },
            }
            for t in selected
        ]

    def to_anthropic_format(self, names: list[str] | None = None) -> list[dict]:
        """Translate registry entries to Anthropic tool format."""
        selected = self._select(names)
        return [
            {
                "name": t.name,
                "description": t.description,
                "input_schema": t.schema,
            }
            for t in selected
        ]

    def to_nexus_format(self, names: list[str] | None = None) -> list[ToolDefinition]:
        """Translate registry entries to canonical NEXUS definitions."""
        selected = self._select(names)
        return [ToolDefinition(name=t.name, description=t.description, parameters=t.schema) for t in selected]

    def _select(self, names: list[str] | None) -> list[ToolEntry]:
        if names is None:
            return list(self._tools.values())
        return [self._tools[n] for n in names if n in self._tools]

    def _schema_from_callable(self, fn) -> dict[str, Any]:
        sig = inspect.signature(fn)
        hints = get_type_hints(fn)
        properties: dict[str, Any] = {}
        required: list[str] = []

        for pname, param in sig.parameters.items():
            ann = hints.get(pname, Any)
            optional = self._is_optional(ann)
            unwrapped = self._unwrap_optional(ann)
            properties[pname] = self._annotation_to_schema(unwrapped)
            if not optional and param.default is inspect._empty:
                required.append(pname)

        return {
            "type": "object",
            "properties": properties,
            "required": required,
            "additionalProperties": False,
        }

    def _annotation_to_schema(self, ann: Any) -> dict[str, Any]:
        origin = get_origin(ann)
        if ann is str:
            return {"type": "string"}
        if ann is int:
            return {"type": "integer"}
        if ann is float:
            return {"type": "number"}
        if ann is bool:
            return {"type": "boolean"}
        if ann in {list, tuple} or origin in {list, tuple}:
            return {"type": "array"}
        if ann is dict or origin is dict:
            return {"type": "object"}
        return {"type": "string"}

    def _is_optional(self, ann: Any) -> bool:
        origin = get_origin(ann)
        if origin in {Union, getattr(__import__("types"), "UnionType", Union)}:
            return NoneType in get_args(ann)
        return False

    def _unwrap_optional(self, ann: Any) -> Any:
        if not self._is_optional(ann):
            return ann
        args = [a for a in get_args(ann) if a is not NoneType]
        return args[0] if args else str
