# -*- coding: utf-8 -*-

"""
Copyright 2025 The Dapr 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
    http://www.apache.org/licenses/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 __future__ import annotations

import asyncio
import inspect
import json
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Mapping, Optional, Protocol, Sequence, Union, cast

from google.protobuf.any_pb2 import Any as GrpcAny

from dapr.clients.grpc import _conversation_helpers as conv_helpers
from dapr.clients.grpc._conversation_helpers import _generate_unique_tool_call_id
from dapr.proto import api_v1

Params = Union[Mapping[str, Any], Sequence[Any], None]

# ------------------------------------------------------------------------------------------------
# Request Classes
# ------------------------------------------------------------------------------------------------


@dataclass
class ConversationInput:
    """A single input message for the conversation."""

    content: str
    role: Optional[str] = None
    scrub_pii: Optional[bool] = None


@dataclass
class ConversationMessageContent:
    """Content for conversation messages."""

    text: str


def _indent_lines(title: str, text: str, indent: int) -> str:
    """
    Indent lines of text.
    Example:
    >>> print("foo")
    foo
    >>> print(_indent_lines("Description", "This is a long\nmultiline\ntext block", 4))
        Description: This is a long
                     multiline
                     text block
    """
    indent_after_first_line = indent + len(title) + 2
    lines = text.splitlines() if text is not None else ['']
    first = lines[0] if lines else ''
    rest = ''
    for line in lines[1:]:
        rest += f'\n{indent_after_first_line * " "}{line}'
    return f'{indent * " "}{title}: {first}{rest}'


class HasNameAndContent:
    """Mixin Protocol for name and content typing."""

    name: Optional[str] = None
    content: List[ConversationMessageContent] = field(default_factory=list)


class UserTracePrintMixin(HasNameAndContent):
    """Mixin for trace_print for text based message content from user to LLM."""

    def trace_print(self, indent: int = 0) -> None:
        base = ' ' * indent
        if self.name:
            print(f'{base}name: {self.name}')
        for i, c in enumerate(self.content):
            print(_indent_lines(f'content[{i}]', c.text, indent))


@dataclass
class ConversationMessageOfDeveloper(UserTracePrintMixin):
    """Developer message content."""

    name: Optional[str] = None
    content: List[ConversationMessageContent] = field(default_factory=list)


@dataclass
class ConversationMessageOfSystem(UserTracePrintMixin):
    """System message content."""

    name: Optional[str] = None
    content: List[ConversationMessageContent] = field(default_factory=list)


@dataclass
class ConversationMessageOfUser(UserTracePrintMixin):
    """User message content."""

    name: Optional[str] = None
    content: List[ConversationMessageContent] = field(default_factory=list)


@dataclass
class ConversationToolCallsOfFunction:
    """Function call details within a tool call."""

    name: str
    arguments: str


@dataclass
class ConversationToolCalls:
    """Tool calls generated by the model."""

    id: Optional[str] = None
    function: Optional[ConversationToolCallsOfFunction] = None


@dataclass
class ConversationMessageOfAssistant:
    """Assistant message content."""

    name: Optional[str] = None
    content: List[ConversationMessageContent] = field(default_factory=list)
    tool_calls: List[ConversationToolCalls] = field(default_factory=list)

    def trace_print(self, indent: int = 0) -> None:
        base = ' ' * indent
        if self.name:
            print(f'{base}name: {self.name}')
        for i, c in enumerate(self.content):
            print(_indent_lines(f'content[{i}]', c.text, indent))
        if self.tool_calls:
            print(f'{base}tool_calls: {len(self.tool_calls)}')
            for idx, tc in enumerate(self.tool_calls):
                tc_id = tc.id or ''
                fn = tc.function.name if tc.function else ''
                args = tc.function.arguments if tc.function else ''
                print(f'{base}  [{idx}] id={tc_id} function={fn}({args})')


@dataclass
class ConversationMessageOfTool:
    """Tool message content."""

    tool_id: Optional[str] = None
    name: str = ''
    content: List[ConversationMessageContent] = field(default_factory=list)

    def trace_print(self, indent: int = 0) -> None:
        base = ' ' * indent
        if self.tool_id:
            print(f'{base}tool_id: {self.tool_id}')
        if self.name:
            print(f'{base}name: {self.name}')
        for i, c in enumerate(self.content):
            lines = c.text.splitlines() if c.text is not None else ['']
            first = lines[0] if lines else ''
            print(f'{base}content[{i}]: {first}')
            for extra in lines[1:]:
                print(extra)


@dataclass
class ConversationMessage:
    """Conversation message with different role types."""

    of_developer: Optional[ConversationMessageOfDeveloper] = None
    of_system: Optional[ConversationMessageOfSystem] = None
    of_user: Optional[ConversationMessageOfUser] = None
    of_assistant: Optional[ConversationMessageOfAssistant] = None
    of_tool: Optional[ConversationMessageOfTool] = None

    def trace_print(self, indent: int = 0):
        print()
        """Print the conversation message with indentation and direction arrows."""
        if self.of_developer:
            print(f'{" " * indent}client[devel]  --------------> LLM[assistant]:')
            self.of_developer.trace_print(indent + 2)
        if self.of_system:
            print(f'{" " * indent}client[system] --------------> LLM[assistant]:')
            self.of_system.trace_print(indent + 2)
        if self.of_user:
            print(f'{" " * indent}client[user]   --------------> LLM[assistant]:')
            self.of_user.trace_print(indent + 2)
        if self.of_assistant:
            print(f'{" " * indent}client         <------------- LLM[assistant]:')
            self.of_assistant.trace_print(indent + 2)
        if self.of_tool:
            print(f'{" " * indent}client[tool]   -------------> LLM[assistant]:')
            self.of_tool.trace_print(indent + 2)

    def to_proto(self) -> api_v1.ConversationMessage:
        """Convert a conversation message to proto format."""

        def _convert_message_content_to_proto(
            content_list: List[ConversationMessageContent],
        ):
            """Convert message content list to proto format."""
            if not content_list:
                return []
            return [
                api_v1.ConversationMessageContent(text=content.text) for content in content_list
            ]

        def _convert_tool_calls_to_proto(tool_calls: List[ConversationToolCalls]):
            """Convert tool calls to proto format."""
            if not tool_calls:
                return []
            proto_calls = []
            for call in tool_calls:
                proto_call = api_v1.ConversationToolCalls()
                if call.id:
                    proto_call.id = call.id
                if call.function:
                    proto_call.function.name = call.function.name
                    proto_call.function.arguments = call.function.arguments
                proto_calls.append(proto_call)
            return proto_calls

        proto_message = api_v1.ConversationMessage()

        if self.of_developer:
            proto_message.of_developer.name = self.of_developer.name or ''
            proto_message.of_developer.content.extend(
                _convert_message_content_to_proto(self.of_developer.content or [])
            )
        elif self.of_system:
            proto_message.of_system.name = self.of_system.name or ''
            proto_message.of_system.content.extend(
                _convert_message_content_to_proto(self.of_system.content or [])
            )
        elif self.of_user:
            proto_message.of_user.name = self.of_user.name or ''
            proto_message.of_user.content.extend(
                _convert_message_content_to_proto(self.of_user.content or [])
            )
        elif self.of_assistant:
            proto_message.of_assistant.name = self.of_assistant.name or ''
            proto_message.of_assistant.content.extend(
                _convert_message_content_to_proto(self.of_assistant.content or [])
            )
            proto_message.of_assistant.tool_calls.extend(
                _convert_tool_calls_to_proto(self.of_assistant.tool_calls or [])
            )
        elif self.of_tool:
            if self.of_tool.tool_id:
                proto_message.of_tool.tool_id = self.of_tool.tool_id
            proto_message.of_tool.name = self.of_tool.name
            proto_message.of_tool.content.extend(
                _convert_message_content_to_proto(self.of_tool.content or [])
            )

        return proto_message


@dataclass
class ConversationInputAlpha2:
    """Alpha2 input message for conversation API."""

    messages: List[ConversationMessage]
    scrub_pii: Optional[bool] = None


@dataclass
class ConversationToolsFunction:
    """Function definition for conversation tools."""

    name: str
    description: Optional[str] = None
    parameters: Optional[Dict] = None

    def schema_as_dict(self) -> Dict:
        """Return the function's schema as a dictionary.

        Returns:
            Dict: The JSON schema for the function parameters.
        """
        return self.parameters or {}

    @classmethod
    def from_function(cls, func: Callable, register: bool = True) -> 'ConversationToolsFunction':
        """Create a ConversationToolsFunction from a function.

        Args:
            func: The function to extract the schema from.
            register: Whether to register the function in the tool registry.
        """
        c = cls(
            name=func.__name__,
            description=conv_helpers.extract_docstring_summary(func),
            parameters=conv_helpers.function_to_json_schema(func),
        )
        if register:
            register_tool(c.name, ConversationTools(function=c, backend=FunctionBackend(func)))
        return c


# ------------------------------------------------------------------------------------------------
# Response Classes
# ------------------------------------------------------------------------------------------------


@dataclass
class ConversationResultAlpha1:
    """One of the outputs to a request to the conversation API."""

    result: str
    parameters: Dict[str, GrpcAny] = field(default_factory=dict)


@dataclass
class ConversationResultAlpha2Message:
    """Message content in one conversation result choice."""

    content: str
    tool_calls: List[ConversationToolCalls] = field(default_factory=list)


@dataclass
class ConversationResultAlpha2Choices:
    """Choice in one Alpha2 conversation result output."""

    finish_reason: str
    index: int
    message: ConversationResultAlpha2Message


@dataclass
class ConversationResultAlpha2:
    """One of the outputs in Alpha2 response from conversation input."""

    choices: List[ConversationResultAlpha2Choices] = field(default_factory=list)


@dataclass
class ConversationResponseAlpha1:
    """Response to a request from the conversation API."""

    context_id: Optional[str]
    outputs: List[ConversationResultAlpha1]


@dataclass
class ConversationResponseAlpha2:
    """Alpha2 response to a request from the conversation API."""

    context_id: Optional[str]
    outputs: List[ConversationResultAlpha2]

    def to_assistant_messages(self) -> List[ConversationMessage]:
        """Helper to convert to Assistant messages and makes it easy to use in multi-turn conversations."""

        def convert_llm_response_to_conversation_input(
            result_message: ConversationResultAlpha2Message,
        ) -> ConversationMessage:
            """Convert ConversationResultMessage (from LLM response) to ConversationMessage."""

            # Convert content string to ConversationMessageContent list
            content = []
            if result_message.content:
                content = [ConversationMessageContent(text=(result_message.content))]

            # Convert tool_calls if present (they're already the right type)
            tool_calls = result_message.tool_calls or []

            # Create an assistant message (since LLM responses are always assistant messages)
            return ConversationMessage(
                of_assistant=ConversationMessageOfAssistant(content=content, tool_calls=tool_calls)
            )

        """Convert the outputs to a list of ConversationInput."""
        assistant_messages = []

        for output in self.outputs or []:
            for choice in output.choices or []:
                # Convert and collect all assistant messages
                assistant_message = convert_llm_response_to_conversation_input(choice.message)
                assistant_messages.append(assistant_message)

        return assistant_messages


# ------------------------------------------------------------------------------------------------
# Tool Helpers
# ------------------------------------------------------------------------------------------------


class ToolBackend(Protocol):
    """Interface for executors that knows how to execute a tool call."""

    def invoke(self, spec: ConversationToolsFunction, params: Params) -> Any: ...

    async def ainvoke(
        self, spec: ConversationToolsFunction, params: Params, *, timeout: Union[float, None] = None
    ) -> Any: ...


@dataclass
class FunctionBackend:
    """A backend that executes a local function."""

    func: Callable[..., Any] = field(repr=False)

    def invoke(self, spec: ConversationToolsFunction, params: Params) -> Any:
        bound = conv_helpers.bind_params_to_func(self.func, params)
        if inspect.iscoroutinefunction(self.func):
            raise conv_helpers.ToolExecutionError(
                "This tool is async; use 'await tool.ainvoke(...)'."
            )
        try:
            return self.func(*bound.args, **bound.kwargs)
        except Exception as e:
            raise conv_helpers.ToolExecutionError(f'Tool raised: {e}') from e

    async def ainvoke(
        self, spec: ConversationToolsFunction, params: Params, *, timeout: Union[float, None] = None
    ) -> Any:
        bound = conv_helpers.bind_params_to_func(self.func, params)
        try:
            if inspect.iscoroutinefunction(self.func):
                coro = self.func(*bound.args, **bound.kwargs)
                return await (asyncio.wait_for(coro, timeout) if timeout else coro)
            loop = asyncio.get_running_loop()
            return await loop.run_in_executor(None, lambda: self.func(*bound.args, **bound.kwargs))
        except asyncio.TimeoutError as err:
            raise conv_helpers.ToolExecutionError(f'Timed out after {timeout} seconds') from err
        except Exception as e:
            raise conv_helpers.ToolExecutionError(f'Tool raised: {e}') from e


def tool(
    func: Optional[Callable] = None,
    *,
    name: Optional[str] = None,
    description: Optional[str] = None,
    namespace: Optional[str] = None,
    register: bool = True,
):
    """
    Decorate a callable as a conversation tool.

    Security note:
    - Register only trusted functions. Tool calls may be triggered from LLM outputs and receive
      untrusted parameters.
    - Use precise type annotations and docstrings for your function; we derive a JSON schema used by
      the binder to coerce types and reject unexpected/invalid arguments.
    - Add your own guardrails if the tool can perform side effects (filesystem, network, subprocess).
    - You can set register=False and call register_tool later to control registration explicitly.
    """

    def _decorate(f: Callable):
        ctf = ConversationToolsFunction.from_function(f, register=False)

        # Prefix name with namespace/module if not provided explicitly
        ns = namespace or ''
        if ns:
            ns += '.'
        ctf.name = name or f'{ns}{ctf.name}'

        if description:
            ctf.description = description

        ct = ConversationTools(function=ctf, backend=FunctionBackend(f))

        # Store the tool in the function for later retrieval (mypy-safe without setattr)
        cast(Any, f).__dapr_conversation_tool__ = ct

        if register:
            register_tool(ctf.name, ct)

        return f

    return _decorate if func is None else _decorate(func)


@dataclass
class ConversationTools:
    """Tools available for conversation.

    Notes on safety and validation:
    - Tools execute arbitrary Python callables. Register only trusted functions and be mindful of
      side effects (filesystem, network, subprocesses).
    - Parameters provided by an LLM are untrusted. The invocation path uses bind_params_to_func to
      coerce types based on your function annotations and to reject unexpected/invalid arguments.
    - Consider adding your own validation/guardrails in your tool implementation.
    """

    # currently only function is supported
    function: ConversationToolsFunction
    backend: Optional[ToolBackend] = None

    def invoke(self, params: Params = None) -> Any:
        """Execute the tool with params (synchronous).

        params may be:
        - Mapping[str, Any]: passed as keyword arguments
        - Sequence[Any]: passed as positional arguments
        - None: no arguments
        Detailed validation and coercion are performed by the backend via bind_params_to_func.
        """
        if not self.backend:
            raise conv_helpers.ToolExecutionError('Tool backend not set')
        return self.backend.invoke(self.function, params)

    async def ainvoke(self, params: Params = None, *, timeout: Union[float, None] = None) -> Any:
        """Execute the tool asynchronously. See invoke() for parameter shape and safety notes."""
        if not self.backend:
            raise conv_helpers.ToolExecutionError('Tool backend not set')
        return await self.backend.ainvoke(self.function, params, timeout=timeout)


# registry of tools
_TOOL_REGISTRY: Dict[str, ConversationTools] = {}


def register_tool(name: str, t: ConversationTools):
    if name in _TOOL_REGISTRY:
        raise ValueError(f"Tool '{name}' already registered")
    _TOOL_REGISTRY[name] = t


def unregister_tool(name: str):
    """Unregister a tool. Good for cleanup and avoid collisions."""
    if name in _TOOL_REGISTRY:
        del _TOOL_REGISTRY[name]


def get_registered_tools() -> List[ConversationTools]:
    """Get a list of all registered tools. This can be pass as tools in the ConversationInput."""
    return list(_TOOL_REGISTRY.values())


def _get_tool(name: str) -> ConversationTools:
    try:
        return _TOOL_REGISTRY[name]
    except KeyError as err:
        raise conv_helpers.ToolNotFoundError(f"Tool '{name}' is not registered") from err


def execute_registered_tool(name: str, params: Union[Params, str] = None) -> Any:
    """Execute a registered tool.

    Security considerations:
    - A registered tool typically executes user-defined code (or code imported from libraries). Only
      register and execute tools you trust. Treat model-provided params as untrusted input.
    - Prefer defining a JSON schema for your tool function parameters (ConversationToolsFunction
      is created from your function’s signature and annotations). The internal binder performs
      type coercion and rejects unexpected/invalid arguments.
    - Add your own guardrails if the tool can perform side effects (filesystem, network, subprocess, etc.).
    """
    if isinstance(params, str):
        params = json.loads(params)
    # Minimal upfront shape check; detailed validation happens in bind_params_to_func
    if params is not None and not isinstance(params, (Mapping, Sequence)):
        raise conv_helpers.ToolArgumentError(
            'params must be a mapping (kwargs), a sequence (args), or None'
        )
    return _get_tool(name).invoke(params)


async def execute_registered_tool_async(
    name: str, params: Union[Params, str] = None, *, timeout: Union[float, None] = None
) -> Any:
    """Execute a registered tool asynchronously.

    Security considerations:
    - Only execute trusted tools; treat model-provided params as untrusted input.
    - Prefer well-typed function signatures and schemas for parameter validation. The binder will
      coerce and validate, rejecting unexpected arguments.
    - For async tools, consider timeouts and guardrails to limit side effects.
    """
    if isinstance(params, str):
        params = json.loads(params)
    if params is not None and not isinstance(params, (Mapping, Sequence)):
        raise conv_helpers.ToolArgumentError(
            'params must be a mapping (kwargs), a sequence (args), or None'
        )
    return await _get_tool(name).ainvoke(params, timeout=timeout)


# ------------------------------------------------------------------------------------------------
# Helpers to create messages for Alpha2 inputs
# ------------------------------------------------------------------------------------------------


def create_user_message(text: str) -> ConversationMessage:
    """Helper to create a user message for Alpha2."""
    return ConversationMessage(
        of_user=ConversationMessageOfUser(content=[ConversationMessageContent(text=text)])
    )


def create_system_message(text: str) -> ConversationMessage:
    """Helper to create a system message for Alpha2."""
    return ConversationMessage(
        of_system=ConversationMessageOfSystem(content=[ConversationMessageContent(text=text)])
    )


def create_assistant_message(text: str) -> ConversationMessage:
    """Helper to create an assistant message for Alpha2."""
    return ConversationMessage(
        of_assistant=ConversationMessageOfAssistant(content=[ConversationMessageContent(text=text)])
    )


def create_tool_message(tool_id: str, name: str, content: Any) -> ConversationMessage:
    """Helper to create a tool message for Alpha2 responses (from client to LLM)."""
    content = conv_helpers.stringify_tool_output(content)
    return ConversationMessage(
        of_tool=ConversationMessageOfTool(
            tool_id=tool_id, name=name, content=[ConversationMessageContent(text=content)]
        )
    )


def _get_outputs_from_grpc_response(
    response: api_v1.ConversationResponseAlpha2,
) -> List[ConversationResultAlpha2]:
    """Helper to get outputs from a Converse gRPC response from dapr sidecar."""
    outputs: List[ConversationResultAlpha2] = []
    for output in response.outputs:
        choices = []
        for choice in output.choices:
            # Convert tool calls from response
            tool_calls = []
            for tool_call in choice.message.tool_calls:
                function_call = ConversationToolCallsOfFunction(
                    name=tool_call.function.name, arguments=tool_call.function.arguments
                )
                if not tool_call.id:
                    tool_call.id = _generate_unique_tool_call_id()
                tool_calls.append(ConversationToolCalls(id=tool_call.id, function=function_call))

            result_message = ConversationResultAlpha2Message(
                content=choice.message.content, tool_calls=tool_calls
            )

            choices.append(
                ConversationResultAlpha2Choices(
                    finish_reason=choice.finish_reason,
                    index=choice.index,
                    message=result_message,
                )
            )

        outputs.append(ConversationResultAlpha2(choices=choices))
    return outputs
