"""
The ``queries`` module provides abstractions for managing pre-defined RAI queries
that Cortex agents can discover and execute.
"""
import inspect
from abc import ABC
from typing import Callable, Dict, Any, Optional, Union
import warnings

import pandas
import relationalai.semantics as rai

from relationalai.util.docutils import include_in_docs

__include_in_docs__ = True


@include_in_docs
class Queries(ABC):
    """Abstract base class for managing pre-defined RAI queries.

    Queries provide a registry of pre-built query functions that Cortex agents
    can discover and execute. This allows domain experts to define common
    analytical queries that agents can use without writing RAI code.

    Implement this interface to create custom query providers, or use
    ``QueryCatalog`` for standard Python function-based queries.
    """

    @include_in_docs
    def get_available_queries(self) -> dict:
        """Return metadata about available queries.

        Returns
        -------
        dict
            Dictionary with the following keys:

            queries : dict
                Mapping from a query ID to a query metadata dictionary.
                Each query metadata dictionary contains:

                id : str
                    Unique query identifier.
                description : str
                    Human-readable description of the query.
                args : dict
                    Mapping from parameter name to parameter type.

            query_instructions : str
                Instructions describing how to use the catalog.
        """
        raise NotImplementedError()

    @include_in_docs
    def explain(self, query_id: str) -> Optional[str]:
        """Return source code for a query by ID, or None if not found.

        Parameters
        ----------
        query_id : str
            Unique identifier for the query to explain.

        Returns
        -------
        str or None
            Source code string, or ``None`` if the query ID is not found.
        """
        raise NotImplementedError()

    @include_in_docs
    def execute(self, query_id: str, args: dict) -> Union[rai.Fragment, pandas.DataFrame]:
        """Execute a specific query by ID with provided arguments.

        Parameters
        ----------
        query_id : str
            Unique identifier for the query to execute (e.g.,
            ``"top_customers"``).
        args : dict
            Dictionary mapping parameter names to values (e.g.,
            ``{"limit": 5, "days": 30}``).

        Returns
        -------
        Union[rai.Fragment, pandas.DataFrame]
            Query results.

        Raises
        ------
        KeyError
            If the ``query_id`` is not found in registered queries.
        """
        raise NotImplementedError()


@include_in_docs
class QueryCatalog(Queries):
    """Query provider that registers Python functions as executable queries.

    Extracts metadata from Python callables (functions, methods) including
    their docstrings, type annotations, and source code. This metadata is
    exposed to Cortex agents for query discovery, while the functions remain
    executable for query execution.

    Query functions should follow these conventions:

    1. Return type: Must return either ``rai.Fragment`` or ``pandas.DataFrame``
    2. Parameters: Can have any parameters with type annotations (recommended)
    3. Docstring: Should include a clear description of what the query does
    4. Function name: Will be used as the query ID (e.g., ``"top_customers"``)

    Parameters
    ----------
    *query_funcs : Callable
        Variable number of Python callables to register as queries.

    Examples
    --------

    >>> from functools import partial
    >>> def top_customers(model, limit=10):
    ...     return model.query(...)
    >>> queries = QueryCatalog(partial(top_customers, jaffle_model))
    """

    def __init__(self, *query_funcs: Callable):
        self._queries: Dict[str, Dict[str, Any]] = {}

        for query_func in query_funcs:
            try:
                # Extract full metadata including source code
                metadata = _get_query_metadata(query_func)
                self._queries[metadata["id"]] = metadata
            except (ValueError, TypeError, AttributeError) as e:
                # If full metadata extraction fails, register with minimal info
                # This ensures queries remain executable even in restricted environments
                warnings.warn(
                    f"Could not extract full metadata for query function "
                    f"{getattr(query_func, '__name__', 'unknown')}: {e}. "
                    f"Registering with minimal metadata.",
                    UserWarning
                )
                fallback_metadata = _get_fallback_metadata(query_func, str(e))
                self._queries[fallback_metadata["id"]] = fallback_metadata

    @include_in_docs
    def get_available_queries(self) -> dict:
        """Return metadata for all registered queries.

        Returns
        -------
        dict
            Dict with ``"queries"`` and ``"query_instructions"`` keys.
        """
        # Filter out 'f' (non-serializable function reference) and 'code'
        # (verbose implementation detail unnecessary for discovery)
        queries_metadata = {
            query_id: {key: value for key, value in metadata.items() if key not in ("f", "code")}
            for query_id, metadata in self._queries.items()
        }

        return {
            "queries": queries_metadata,
            "query_instructions": (
                "IMPORTANT: Only these specific pre-defined queries are available. "
                "You CANNOT create custom queries or answer arbitrary analytical questions. "
                "For questions matching these pre-defined queries, provide a JSON payload with 'id' and 'args'. "
                "For other analytical questions, clearly explain you can only run these specific pre-defined queries."
            ),
        }

    @include_in_docs
    def explain(self, query_id: str) -> Optional[str]:
        """Return source code for a registered query by ID.

        Parameters
        ----------
        query_id : str
            Unique identifier for the query to explain.

        Returns
        -------
        str or None
            Source code string, or ``None`` if the query ID is not found.
        """
        if query_id not in self._queries:
            return None
        if "code" not in self._queries[query_id]:
            return None
        return self._queries[query_id]["code"]

    @include_in_docs
    def execute(self, query_id: str, args: dict) -> Union[rai.Fragment, pandas.DataFrame]:
        """Execute a registered query by ID.

        Parameters
        ----------
        query_id : str
            ID of the query to execute.
        args : dict
            Arguments to pass to the query function.

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

        Raises
        ------
        KeyError
            If ``query_id`` is not registered.
        """
        if query_id not in self._queries:
            available = list(self._queries.keys())
            raise KeyError(
                f"Query '{query_id}' not found. "
                f"This model has a limited catalog with only these pre-defined queries: {available}. "
                f"Custom queries cannot be executed - only these specific analytical queries are available."
            )

        query_func = self._queries[query_id]["f"]
        return query_func(**args)


def _get_query_metadata(func: Callable) -> Dict[str, Any]:
    """Extract complete metadata from a query function.

    Parameters
    ----------
    func : Callable
        The query function to extract metadata from.

    Returns
    -------
    dict
        Dict containing ``'id'``, ``'description'``, ``'args'``, ``'code'``,
        and ``'f'`` (function reference).

    Raises
    ------
    ValueError
        If function name cannot be determined.
    TypeError
        If signature cannot be extracted.
    """
    # Get function signature
    sig = inspect.signature(func)

    # Get function name
    if not hasattr(func, '__name__'):
        raise ValueError(f"Function {func} does not have a __name__ attribute")
    name = func.__name__

    # Get docstring as description
    description = inspect.getdoc(func) or ""

    # Extract arguments and their types
    args = {}
    for param_name, param in sig.parameters.items():
        if param.annotation != inspect.Parameter.empty:
            # Get string representation of the type annotation
            if hasattr(param.annotation, '__name__'):
                args[param_name] = param.annotation.__name__
            else:
                args[param_name] = str(param.annotation)
        else:
            args[param_name] = "Any"

    # Get source code - handle bound methods and other edge cases
    try:
        # For bound methods, get the underlying function
        if hasattr(func, '__func__'):
            code = inspect.getsource(getattr(func, '__func__'))
        else:
            code = inspect.getsource(func)
    except (OSError, TypeError, AttributeError):
        # If we can't get source, provide a placeholder
        code = f"# Source code for {name} (unable to retrieve in this environment)\n"

    return {
        "id": name,
        "description": description,
        "args": args,
        "code": code,
        "f": func,  # Keep function reference for execution
    }


def _get_fallback_metadata(func: Callable, error_msg: str) -> Dict[str, Any]:
    """Create minimal metadata when full extraction fails.

    Parameters
    ----------
    func : Callable
        The query function.
    error_msg : str
        Error message from the failed extraction.

    Returns
    -------
    dict
        Minimal metadata that still allows query execution.
    """
    # Try to get function name, fallback to 'unknown'
    func_name = getattr(func, '__name__', 'unknown')

    # Try to get docstring even if other metadata extraction failed
    try:
        description = inspect.getdoc(func) or ""
    except Exception:
        description = ""

    # Try to extract arguments even if source code failed
    args = {}
    try:
        sig = inspect.signature(func)
        for param_name, param in sig.parameters.items():
            if param.annotation != inspect.Parameter.empty:
                if hasattr(param.annotation, '__name__'):
                    args[param_name] = param.annotation.__name__
                else:
                    args[param_name] = str(param.annotation)
            else:
                args[param_name] = "Any"
    except Exception:
        # If signature extraction also fails, leave args empty
        pass

    return {
        "id": func_name,
        "description": description,
        "args": args,
        "code": f"# Query {func_name} (metadata extraction failed: {error_msg})\n",
        "f": func,  # Critical: keep function reference for execution
    }
