"""
The ``tool`` module provides the tool abstraction and registry for RAI-powered
Cortex agent tools.
"""
from abc import ABC, abstractmethod
from enum import Enum
from typing import List, Any, Union, Optional, Set

import pandas
import relationalai.semantics as rai

from relationalai.util.docutils import include_in_docs

from .verbalize import Verbalizer, ModelVerbalizer
from .queries import Queries

__include_in_docs__ = True


@include_in_docs
class MaturityLevel(Enum):
    """Capability maturity levels for RAI tools.

    Attributes
    ----------
    GA : str
        Generally Available.
    PREVIEW : str
        Preview capability with potential limitations.
    """
    GA = "GA"
    PREVIEW = "Preview"


class Capability(Enum):
    """Tool capabilities that can be discovered and used by agents."""
    DISCOVER = ("discover", MaturityLevel.GA)
    VERBALIZE = ("verbalize", MaturityLevel.GA)
    EXPLAIN = ("explain", MaturityLevel.GA)
    QUERY = ("query", MaturityLevel.PREVIEW)

    def __init__(self, value: str, maturity: MaturityLevel):
        self._value_ = value
        self.maturity = maturity

    @property
    def is_preview(self) -> bool:
        return self.maturity == MaturityLevel.PREVIEW


@include_in_docs
class Tool(ABC):
    """Abstract base class for RAI-powered tools that can be used by Cortex agents.

    A Tool provides semantic data exploration capabilities through a consistent
    interface:

    - Discovery: What models and concepts are available?
    - Verbalization: Detailed model structure and relationships
    - Explanation: Business rules and semantics for specific concepts
    - Querying: Execute structured queries against the data model

    Implement this interface to create custom tools, or use ``DefaultTool`` for
    standard RAI model-based tools.
    """

    @include_in_docs
    @abstractmethod
    def description(self) -> str:
        """Return a human-readable description of this tool's purpose.

        Returns
        -------
        str
            A description of the tool.
        """
        pass

    @include_in_docs
    @abstractmethod
    def concepts(self) -> List[str]:
        """Return list of key concept names available in this tool's domain.

        Returns
        -------
        list of str
            Concept names available in this tool.
        """
        pass

    @include_in_docs
    @abstractmethod
    def verbalize(self) -> str:
        """Return detailed structure and relationships for the entire model.

        Returns
        -------
        str
            A natural language description of the model structure.
        """
        pass

    @include_in_docs
    @abstractmethod
    def explain(self, concept: str) -> str:
        """Return detailed explanation of business rules for a specific concept.

        Parameters
        ----------
        concept : str
            The name of the concept to explain.

        Returns
        -------
        str
            A natural language explanation of the concept.
        """
        pass

    def _capabilities(self) -> Set[Capability]:
        """Return the set of capabilities this tool supports.

        Returns
        -------
        set of Capability
            Capabilities supported by this tool. Base implementation returns
            DISCOVER, VERBALIZE, and EXPLAIN.
        """
        return {Capability.DISCOVER, Capability.VERBALIZE, Capability.EXPLAIN}

    @include_in_docs
    @abstractmethod
    def queries(self) -> dict:
        """Return available pre-defined queries.

        If queries capability is not configured, returns a message
        explaining that no queries are available.

        Returns
        -------
        dict
            Dict with ``"queries"`` and ``"query_instructions"`` if configured,
            or a dict with ``"message"`` explaining queries are not available.
        """
        return {}

    @include_in_docs
    @abstractmethod
    def query(self, params: dict) -> Union[rai.Fragment, pandas.DataFrame]:
        """Execute a query against this tool's model and return results.

        Parameters
        ----------
        params : dict
            Query parameters including ``"id"`` (query identifier) and
            ``"args"`` (query arguments).

        Returns
        -------
        rai.Fragment or pandas.DataFrame
            Query results.

        Raises
        ------
        ValueError
            If queries capability is not configured.
        KeyError
            If query ID is not found.
        """
        raise NotImplementedError()


@include_in_docs
class DefaultTool(Tool):
    """Standard implementation of Tool for RAI semantic models.

    Wraps a RAI Model and provides semantic exploration capabilities through
    discovery, verbalization, explanation, and querying. Uses pluggable
    components (``Verbalizer`` and ``Queries``) to customize how model structure
    and pre-defined queries are presented to the agent.

    Parameters
    ----------
    model : rai.Model
        RAI semantic model to expose through this tool.
    description : str
        Human-readable description of what this tool provides.
    verbalizer : Verbalizer, optional
        Custom verbalizer for model structure. If ``None``, uses
        ``ModelVerbalizer`` which returns relationship readings.
    queries : Queries, optional
        Queries provider for pre-defined query functions (PREVIEW). If
        ``None``, query discovery returns a helpful message.

    Examples
    --------

    >>> tool = DefaultTool(
    ...     model=model,
    ...     description="Jaffle Shop e-commerce data",
    ...     verbalizer=SourceCodeVerbalizer(model, define_customer, define_order),
    ... )
    """

    def __init__(
        self,
        model: rai.Model,
        description: str,
        verbalizer: Optional[Verbalizer] = None,
        queries: Optional[Queries] = None
    ):
        self._model = model
        self._description = description
        self._verbalizer: Verbalizer = verbalizer if verbalizer else ModelVerbalizer(model)
        self._queries: Optional[Queries] = queries

    def description(self) -> str:
        return self._description

    def _capabilities(self) -> Set[Capability]:
        caps = super()._capabilities()
        if self._queries is not None:
            caps = caps | {Capability.QUERY}
        return caps

    def concepts(self) -> List[str]:
        return list(c._name for c in self._model.concepts)

    def verbalize(self) -> str:
        return self._verbalizer.explain_model()

    def explain(self, concept: str) -> str:
        if self._queries is not None:
            query_explanation = self._queries.explain(concept)
            if query_explanation is not None:
                return query_explanation
        return self._verbalizer.explain_concept(concept)

    def queries(self) -> dict:
        if self._queries is None:
            return {
                "message": "No pre-defined queries are configured for this tool. "
                          "To enable queries, provide a Queries instance when creating the tool."
            }
        return self._queries.get_available_queries()

    def query(self, params: dict) -> Union[rai.Fragment, pandas.DataFrame]:
        if self._queries is None:
            raise ValueError(
                "Queries capability not configured. "
                "Provide a Queries instance when creating this tool to enable querying."
            )

        if "id" not in params:
            available = self._queries.get_available_queries()
            available_ids = list(available.get("queries", {}).keys())
            raise ValueError(
                f"Query params must include 'id' key. "
                f'Expected format: {{"id": "<query_id>", "args": {{<params>}}}}. '
                f"Available query IDs: {available_ids}"
            )

        query_id = params["id"]
        args = params.get("args", {})
        return self._queries.execute(query_id, args)


@include_in_docs
class ToolRegistry:
    """Registry for RAI tools accessible to a Cortex agent.

    Provides a unified interface for tool discovery, verbalization,
    explanation, and querying over a RAI model.

    Examples
    --------

    >>> registry = ToolRegistry()
    >>> registry.add(model=model, description="Jaffle shop data")
    """

    def __init__(self, tools=None):
        self._tools: List[Tool] = [] if tools is None else tools

    @include_in_docs
    def add(
        self,
        model: Optional[rai.Model] = None,
        description: Optional[str] = None,
        verbalizer: Optional[Verbalizer] = None,
        queries: Optional[Queries] = None,
        tool: Optional[Tool] = None
    ) -> "ToolRegistry":
        """Add a tool to the registry.

        This method supports two calling patterns:

        1. Provide a ``Tool`` instance directly via the ``tool`` parameter.
        2. Provide ``model`` + ``description`` (creates ``DefaultTool`` internally
           with optional ``verbalizer`` and ``queries``).

        Parameters
        ----------
        model : rai.Model, optional
            RAI Model instance (required if ``tool`` is ``None``).
        description : str, optional
            Human-readable tool description (required if ``tool`` is ``None``).
        verbalizer : Verbalizer, optional
            Custom verbalizer (only used if ``tool`` is ``None``).
        queries : Queries, optional
            Queries provider (only used if ``tool`` is ``None``).
        tool : Tool, optional
            Pre-constructed ``Tool`` instance. If provided, other parameters
            are ignored.

        Returns
        -------
        ToolRegistry
            Self.

        Examples
        --------

        >>> registry = ToolRegistry()
        >>> registry.add(model=model, description="Jaffle shop")
        """
        if tool is not None:
            self._tools.append(tool)
        else:
            if model is None or description is None:
                raise ValueError("Either 'tool' or both 'model' and 'description' must be provided")
            assert model is not None and description is not None  # Type narrowing for pyright
            self._tools.append(DefaultTool(
                model=model,
                description=description,
                verbalizer=verbalizer,
                queries=queries
            ))
        return self

    def __contains__(self, item):
        return 0 <= item < len(self._tools)

    @include_in_docs
    def has_tools(self) -> bool:
        """Check if at least one tool is registered.

        Returns
        -------
        bool
            ``True`` if one or more tools are registered.
        """
        return len(self._tools) > 0

    @include_in_docs
    def has_capability(self, capability: Capability) -> bool:
        """Check if any tool in the registry has the specified capability.

        Parameters
        ----------
        capability : Capability
            The capability to check for.

        Returns
        -------
        bool
            ``True`` if any registered tool has the capability.
        """
        return any(
            self._check_tool_capability(tool, capability)
            for tool in self._tools
        )

    @include_in_docs
    def tool_has_capability(self, tool_id: int, capability: Capability) -> bool:
        """Check if a specific tool has the specified capability.

        Parameters
        ----------
        tool_id : int
            The numeric tool index to check.
        capability : Capability
            The capability to check for.

        Returns
        -------
        bool
            ``True`` if the tool exists and has the capability.
        """
        if not (0 <= tool_id < len(self._tools)):
            return False
        return self._check_tool_capability(self._tools[tool_id], capability)

    @staticmethod
    def _check_tool_capability(tool: Tool, capability: Capability) -> bool:
        """Check if a single tool has a specific capability."""
        return capability in tool._capabilities()

    @include_in_docs
    def discover(self) -> dict:
        """Discover all registered tools and their metadata.

        Returns
        -------
        dict
            Dictionary with ``"models"`` key containing a list of tool metadata
            including model_id, description, key_concepts, and query info.
        """
        return {
            "models": [
                {
                    "model_id": k,
                    "description": tool.description(),
                    "key_concepts": tool.concepts(),
                    **tool.queries(),
                }
                for k, tool in enumerate(self._tools)
            ]
        }

    @include_in_docs
    def verbalize(self, model_id: int) -> str:
        """Get detailed model structure for a specific tool.

        Parameters
        ----------
        model_id : int
            The numeric tool index to verbalize.

        Returns
        -------
        str
            Natural language description of the model structure.
        """
        model = self._tools[model_id]
        return model.verbalize()

    @include_in_docs
    def explain(self, model_id: int, concept: str) -> str:
        """Get detailed explanation of a concept within a specific tool.

        Parameters
        ----------
        model_id : int
            The numeric tool index containing the concept.
        concept : str
            The concept name to explain.

        Returns
        -------
        str
            Natural language explanation of the concept.
        """
        model = self._tools[model_id]
        return model.explain(concept)

    @include_in_docs
    def query(self, model_id: int, query_dict: dict, limit: int = 1000) -> Any:
        """Execute a query against a specific tool.

        Parameters
        ----------
        model_id : int
            The numeric tool index to query.
        query_dict : dict
            Query parameters with ``"id"`` and ``"args"`` keys.
        limit : int
            Maximum number of rows to return. Default: 1000.

        Returns
        -------
        Any
            Query results, typically as a pandas DataFrame.
        """
        model = self._tools[model_id]
        query = model.query(query_dict)
        if isinstance(query, rai.Fragment):
            result = query.to_df().head(limit)
        elif isinstance(query, pandas.DataFrame):
            result = query.head(limit)
        else:
            result = query
        return result
