# Copyright 2025 The RLinf Authors.
#
# Licensed 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
#
#     https://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 dataclasses import dataclass, field
from enum import Enum
from typing import Any, Literal, Optional


class MCPRequestType(Enum):
    """MCP request types."""

    LIST_TOOLS = "list_tools"
    CALL_TOOL = "call_tool"
    LIST_RESOURCES = "list_resources"
    READ_RESOURCE = "read_resource"
    LIST_PROMPTS = "list_prompts"
    GET_PROMPT = "get_prompt"


class MCPSessionState(Enum):
    """MCP session states."""

    INITIALIZING = "initializing"
    CONNECTED = "connected"
    FAILED = "failed"
    TERMINATED = "terminated"


@dataclass
class MCPRequest:
    """
    MCP request structure.
    Used between MCP client(mcp tool worker) and MCP server.
    """

    request_id: str
    request_type: MCPRequestType
    tool_name: Optional[str] = None
    tool_arguments: Optional[dict[str, Any]] = None
    resource_uri: Optional[str] = None
    prompt_name: Optional[str] = None
    prompt_arguments: Optional[dict[str, Any]] = None
    timeout: int = 30
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class MCPResponse:
    """
    MCP response structure.
    Used between MCP client(mcp tool worker) and MCP server.
    """

    request_id: str
    success: bool
    result: Optional[Any] = None
    error_message: Optional[str] = None
    execution_time: float = 0.0
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class ToolChannelRequest:
    """
    Request to a tool channel.
    Used between agent loop worker and tool worker.
    """

    # session_id is REQUIRED. also used as the response channel key
    session_id: str
    request_type: Literal["session_start", "session_end", "execute"]
    tool_name: str | None = None
    tool_args: dict | None = None


@dataclass
class ToolChannelResponse:
    """
    Response from a tool channel.
    Used between agent loop worker and tool worker.
    """

    success: bool

    # execute result if success is True, otherwise failure message
    result: Any | None = None
    session_state: Any | None = None
    meta_info: dict | None = None


@dataclass
class ToolRequest:
    """
    The name of the function to call.
    Used within agent loop worker.
    """

    name: str

    """The arguments to call the function with, as generated by the model in JSON format."""
    arguments: dict


@dataclass
class ToolResponse:
    """
    The response from a tool execution.
    Used within agent loop worker.
    """

    text: str
