################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
from abc import ABC
from typing import Any, Callable, Dict, List, Tuple

from flink_agents.api.function import Function, PythonFunction
from flink_agents.api.resource import (
    ResourceDescriptor,
    ResourceType,
    SerializableResource,
)

STRUCTURED_OUTPUT = "structured_output"


class Agent(ABC):
    """Base class for defining agent logic.


    Example:
        Users have two ways to create an Agent

        * Declare an Agent with decorators
        ::

            class MyAgent(Agent):
                @action(InputEvent.EVENT_TYPE)
                @staticmethod
                def my_action(event: Event, ctx: RunnerContext) -> None:
                    action logic

                @chat_model_connection
                @staticmethod
                def my_connection() -> ResourceDescriptor:
                    return ResourceDescriptor(clazz=OllamaChatModelConnection,
                                              model="qwen2:7b",
                                              base_url="http://localhost:11434")

                @chat_model_setup
                @staticmethod
                def my_chat_model() -> ResourceDescriptor:
                    return ResourceDescriptor(clazz=OllamaChatModel,
                                              connection="my_connection")

        * Add actions and resources to an Agent instance
        ::

            my_agent = Agent()
            my_agent.add_action(name="my_action",
                                events=["_input_event"],
                                func=action_function)
                    .add_resource(name="my_connection",
                                  instance=ResourceDescriptor(
                                        clazz=OllamaChatModelConnection,
                                        arg1=xxx
                                )
                    .add_resource(
                        name="my_connection",
                        instance=ResourceDescriptor(
                            clazz=OllamaChatModelConnection,
                            arg1=xxx
                        )
                    )
                    .add_resource(
                        name="my_chat_model",
                        instance=ResourceDescriptor(
                            clazz=OllamaChatModelSetup,
                            connection="my_connection"
                        )
                    )
    """

    _actions: Dict[
        str, Tuple[List[str], Function, Dict[str, Any] | None]
    ]
    _resources: Dict[ResourceType, Dict[str, Any]]

    def __init__(self) -> None:
        """Init method."""
        self._actions = {}
        self._resources = {}
        for type in ResourceType:
            self._resources[type] = {}

    @property
    def actions(
        self,
    ) -> Dict[str, Tuple[List[str], Function, Dict[str, Any] | None]]:
        """Get added actions."""
        return self._actions

    @property
    def resources(self) -> Dict[ResourceType, Dict[str, Any]]:
        """Get added resources."""
        return self._resources

    def add_action(
        self,
        name: str,
        events: List[str],
        func: Callable | Function,
        **config: Any,
    ) -> "Agent":
        """Add action to agent.

        Parameters
        ----------
        name : str
            The name of the action, should be unique in the same Agent.
        events : list[str]
            Type-identifier strings listened by this action.
        func : Callable | Function
            Either a raw Python callable (it will be wrapped as a
            ``PythonFunction``) or a pre-built flink-agents ``Function``
            (e.g. from the YAML loader).
        **config : Any
            Key named arguments can be used by this action in runtime.

        Returns:
        -------
        Agent
            The modified Agent instance.
        """
        if name in self._actions:
            msg = f"Action {name} already defined"
            raise ValueError(msg)
        if not isinstance(func, Function):
            func = PythonFunction.from_callable(func)
        self._actions[name] = (events, func, config if config else None)
        return self

    def add_resource(
        self,
        name: str,
        resource_type: ResourceType,
        instance: SerializableResource | ResourceDescriptor,
    ) -> "Agent":
        """Add resource to agent instance.

        Parameters
        ----------
        name : str
            The name of the prompt, should be unique in the same Agent.
        resource_type: ResourceType
            The type of the resource.
        instance: SerializableResource | ResourceDescriptor
            The serializable resource instance, or the descriptor of resource.

        Returns:
        -------
        Agent
            The agent to add the resource.
        """
        if name in self._resources[resource_type]:
            msg = f"{resource_type.value} {name} already defined"
            raise ValueError(msg)

        self._resources[resource_type][name] = instance
        return self
