"""
Context.
"""

import textwrap
import warnings
from collections.abc import Callable
from functools import wraps
from inspect import _empty, signature
from typing import (
    Any,
    Literal,
    overload,
)

from pymongo import MongoClient

from ixoncdkingress.function.api_client import ApiClient
from ixoncdkingress.function.document_db_client import (
    DocumentDBClient,
    DocumentType,
)
from ixoncdkingress.function.time_series_client import TimeSeriesClient
from ixoncdkingress.schema_utils import generate_input_schema, generate_output_schema
from ixoncdkingress.types import (
    ObjectJsonSchema,
    Tool,
    ToolIssue,
)


class FunctionResource:
    """
    Describes an IXAPI resource.
    """

    public_id: str
    name: str
    custom_properties: dict[str, Any]
    permissions: set[str] | None

    def __init__(
        self,
        public_id: str,
        name: str,
        custom_properties: dict[str, Any],
        permissions: set[str] | None,
    ) -> None:
        self.public_id = public_id
        self.name = name
        self.custom_properties = custom_properties
        self.permissions = permissions

    def __repr__(self) -> str:
        return (
            "<FunctionResource"
            f" public_id={self.public_id},"
            f" name={self.name},"
            f" custom_properties={self.custom_properties!r},"
            f" permissions={self.permissions!r},"
            f">"
        )


def get_tool_definition(
    function: Callable[..., Any],
    requires_human_permission: bool = False,
    description: str | None = None,
    examples: list[str] | None = None,
    input_schema: ObjectJsonSchema | None = None,
    output_schema: ObjectJsonSchema | None = None,
    meta: dict[str, Any] | None = None,
    exclude_param_types: set[type] | None = None,
    exclude_param_names: set[str] | None = None,
) -> Tool:
    """
    Get the tool definition for a function.
    """
    issues: list[ToolIssue] = []

    _description = function.__doc__ if description is None else description
    if _description is None:
        issues.append(
            ToolIssue(
                type="description_missing",
                error="Agentic methods require description for docstrings",
            )
        )
    else:
        _description = textwrap.dedent(_description).strip(" \n")

    if exclude_param_types is None:
        exclude_param_types = {FunctionContext}
    else:
        exclude_param_types.add(FunctionContext)

    _input_schema = generate_input_schema(
        function,
        exclude_param_types=exclude_param_types,
        exclude_param_names=exclude_param_names,
    )
    _output_schema = generate_output_schema(function)
    if input_schema is not None:
        _input_schema.schema = input_schema
    elif _input_schema.issue is not None:
        issues.append(_input_schema.issue)

    if output_schema is not None:
        _output_schema.schema = output_schema
    elif _output_schema.issue is not None:
        issues.append(_output_schema.issue)

    return Tool(
        name=function.__name__,
        description=_description,
        requires_human_permission=requires_human_permission,
        examples=examples,
        input_schema=_input_schema.schema,
        output_schema=_output_schema.schema,
        meta=meta,
        issues=issues,
        input_model=_input_schema.model,
        output_model=_output_schema.model,
    )


class FunctionContext:
    """
    The context for a Cloud Function.
    """

    config: dict[str, str]
    api_client: ApiClient
    user: FunctionResource | None
    company: FunctionResource | None
    asset: FunctionResource | None
    agent: FunctionResource | None
    template: FunctionResource | None

    _mongo_client: MongoClient[DocumentType] | None
    _ddb_collection_name: str

    @property
    def document_db_client(self) -> DocumentDBClient | None:
        """
        The client for accessing the default collection of the Document DB.

        This may be None when running the Cloud Function in development mode,
        and using the local debug form without specifying a company.
        """
        if not self._mongo_client or not self.company:
            return None

        return DocumentDBClient(
            self._mongo_client[self.company.public_id][self._ddb_collection_name]
        )

    @property
    def time_series_client(self) -> TimeSeriesClient | None:
        """
        The client for accessing the time series collection of the Cloud Function.

        WARNING: This function and client is experimental and may change in the future.

        This may be None when running the Cloud Function in development mode,
        and using the local debug form without specifying a company.
        """
        warnings.warn(
            "The TimeSeriesClient is experimental and may change in the future",
            category=FutureWarning,
        )

        if not self._mongo_client or not self.company:
            return None

        return TimeSeriesClient(
            self._mongo_client[self.company.public_id][f"{self._ddb_collection_name}_ts"]
        )

    @property
    def agent_or_asset(self) -> FunctionResource:
        """
        Return either an Agent or an Asset resource, depending on what's available. If both are
        available in the context, returns the Asset resource.
        """
        if self.asset:
            return self.asset

        assert self.agent
        return self.agent

    def __init__(
        self,
        config: dict[str, str],
        api_client: ApiClient,
        mongo_client: MongoClient[DocumentType] | None,
        user: FunctionResource | None,
        company: FunctionResource | None,
        asset: FunctionResource | None,
        agent: FunctionResource | None,
        template: FunctionResource | None,
        document_db_collection_name: str,
    ) -> None:
        self.config = config
        self.api_client = api_client
        self.user = user
        self.company = company
        self.asset = asset
        self.agent = agent
        self.template = template

        self._mongo_client = mongo_client
        self._ddb_collection_name = document_db_collection_name

    def __repr__(self) -> str:
        return (
            f"<FunctionContext"
            f" config={self.config!r},"
            f" api_client={self.api_client!r},"
            f" document_db_client={self.document_db_client!r},"
            f" user={self.user!r},"
            f" company={self.company!r},"
            f" asset={self.asset!r},"
            f" agent={self.agent!r},"
            f" template={self.template!r},"
            f">"
        )

    @staticmethod
    def expose_cloud_function(function: Any) -> Any:
        """
        Decorator to mark a function as an exposed endpoint.
        """
        sig = signature(function, eval_str=True)

        if not sig.parameters:
            raise Exception("Function has no argument for FunctionContext")

        # If the first function argument has a type annotation it should be of FunctionContext
        context_param = sig.parameters[next(iter(sig.parameters))]
        if (
            context_param.annotation is not _empty
            and context_param.annotation is not FunctionContext
        ):
            raise Exception("First function parameter should be of type FunctionContext")

        function.exposed = True

        return function

    @overload
    @staticmethod
    def expose(
        function: Callable[..., Any],
        *,
        is_agentic: Literal[False] = False,
        requires_human_permission: bool = False,
        description: str | None = None,
        examples: list[str] | None = None,
        input_schema: ObjectJsonSchema | None = None,
        output_schema: ObjectJsonSchema | None = None,
        meta: dict[str, Any] | None = None,
        exclude_param_types: set[type] | None = None,
        exclude_param_names: set[str] | None = None,
    ) -> Callable[..., Any]: ...

    @overload
    @staticmethod
    def expose(  # pragma: no cover
        function: None = None,
        *,
        is_agentic: Literal[True] = True,
        requires_human_permission: bool = False,
        description: str | None = None,
        examples: list[str] | None = None,
        input_schema: ObjectJsonSchema | None = None,
        output_schema: ObjectJsonSchema | None = None,
        meta: dict[str, Any] | None = None,
        exclude_param_types: set[type] | None = None,
        exclude_param_names: set[str] | None = None,
    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: ...

    @staticmethod
    def expose(  # pragma: no cover
        function: Callable[..., Any] | None = None,
        *,
        is_agentic: bool = False,
        requires_human_permission: bool = False,
        description: str | None = None,
        examples: list[str] | None = None,
        input_schema: ObjectJsonSchema | None = None,
        output_schema: ObjectJsonSchema | None = None,
        meta: dict[str, Any] | None = None,
        exclude_param_types: set[type] | None = None,
        exclude_param_names: set[str] | None = None,
    ) -> Callable[[Callable[..., Any]], Callable[..., Any]] | Callable[..., Any]:
        """
        Decorator to mark a function as an exposed tool to an agent.
        Best approach is to use pydantic models for input and output schema, and let this
        decorator generate the schema for you. In case you do not want to use pydantic models,
        and you have non standard python types, you need to provide the input and output
        schema manually.
        """

        def decorator(inner_function: Any) -> Any:
            FunctionContext.expose_cloud_function(inner_function)
            wraps(inner_function)

            if not is_agentic:
                return inner_function

            tool_definition = get_tool_definition(
                function=inner_function,
                description=description,
                examples=examples,
                input_schema=input_schema,
                output_schema=output_schema,
                meta=meta,
                requires_human_permission=requires_human_permission,
                exclude_param_types=exclude_param_types,
                exclude_param_names=exclude_param_names,
            )
            inner_function.tool_definition = tool_definition
            inner_function.exposed_to_agent = is_agentic
            inner_function.input_model = tool_definition.input_model
            inner_function.output_model = tool_definition.output_model

            return inner_function

        return decorator if function is None else decorator(function)
