################################################################################
#  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 typing import Any, ClassVar, Dict, Sequence

from flink_agents.api.agents.agent import Agent
from flink_agents.api.chat_message import ChatMessage
from flink_agents.api.chat_models.chat_model import BaseChatModelSetup
from flink_agents.api.decorators import action, chat_model_setup, tool
from flink_agents.api.events.event import Event, InputEvent
from flink_agents.api.resource import ResourceDescriptor
from flink_agents.api.runner_context import RunnerContext


class MyEvent(Event):
    """Test event."""

    EVENT_TYPE: ClassVar[str] = "_my_event"

    def __init__(self) -> None:
        """Create a MyEvent."""
        super().__init__(type=MyEvent.EVENT_TYPE)


class MockChatModel(BaseChatModelSetup):
    """Mock ChatModel for testing integrating prompt and tool."""

    @property
    def model_kwargs(self) -> Dict[str, Any]:
        """Only for testing."""
        return {}

    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatMessage:
        """Only for test plan compatibility."""


class PythonAgentPlanCompatibilityTestAgent(Agent):
    """Agent for generating python agent plan json."""

    @action(InputEvent.EVENT_TYPE)
    @staticmethod
    def first_action(event: Event, ctx: RunnerContext) -> None:
        """Test implementation."""

    @action("_input_event", "_my_event")
    @staticmethod
    def second_action(event: Event, ctx: RunnerContext) -> None:
        """Test implementation."""

    @chat_model_setup
    @staticmethod
    def chat_model() -> ResourceDescriptor:
        """ChatModel can be used in action."""
        return ResourceDescriptor(
            clazz=f"{MockChatModel.__module__}.{MockChatModel.__name__}",
            name="chat_model",
            prompt="prompt",
            tools=["add"],
            connection="mock_connection",
            model="mock-model",
        )

    @tool
    @staticmethod
    def add(a: int, b: int) -> int:
        """Calculate the sum of a and b.

        Parameters
        ----------
        a : int
            The first operand
        b : int
            The second operand

        Returns:
        -------
        int:
            The sum of a and b
        """
        return a + b
