from typing import Any, Callable, Coroutine, Dict, List, Optional, Type, Union, cast
import asyncio
import logging
from typing import Any, Callable, Dict, Generic, Optional, TypeVar, Union, List
from langchain_core.messages import AIMessage, ToolMessage

import json

from .llm_chat_session import LlmChatSession
from pydantic import PrivateAttr, BaseModel
from .common import deprecated


from langchain_core.tools.base import BaseTool

from .function_decorator import get_function_schema


class ToolsProvider:

    def __init__(
        self,
        tools: Union[
            List[Union[Callable, BaseTool]], Dict[str, Union[Callable, BaseTool]]
        ] = None,
        **kwargs,
    ) -> None:
        """Initialize ToolsProvider with list of tools of dictionary where key is the unique function name alias"""

        if "functions" in kwargs:
            # Backwards compatibility
            tools = kwargs["functions"]
        self.tools = []
        self.aliases = []
        self.tool_schemas = []
        self.tool_name_map = {}
        if not (isinstance(tools, dict) or isinstance(tools, list)):
            raise ValueError(
                "ToolsProvider must be initialized with list of functions or dictionary where key is the unique function name alias"
            )

        for i, f in enumerate(tools):

            if isinstance(f, str):
                function_alias = f
                f = tools[f]
            else:
                function_alias = None

            self.add_tool(f, function_alias)

    def add_tool(self, tool: Union[Callable, BaseTool], alias: str = None):
        """Add function to ToolsProvider. If alias is provided, it will be used as function name in LLM"""
        self.tools.append(tool)
        self.aliases.append(alias)
        if isinstance(tool, BaseTool):
            # self.tool_schemas.append(format_tool_to_openai_function(tool))
            self.tool_schemas.append(tool)
            f_name = alias or tool.name
        elif callable(tool) and hasattr(tool, "get_function_schema"):
            if hasattr(tool, "function_name"):
                f_name = alias or tool.function_name
            else:
                raise Exception(
                    f"Function {tool} does not have function_name attribute. All functions must be marked with @llm_function decorator"
                )
            self.tool_schemas.append(
                lambda kwargs, f=tool: get_function_schema(f, kwargs)
            )
        else:
            raise ValueError(
                f"Invalid item value in functions. Only Tools or functions decorated with @llm_function are allowed. Got: {tool}"
            )
        if f_name in self.tool_name_map:
            if alias:
                raise ValueError(f"Invalid alias - duplicate function name: {f_name}.")
            else:
                raise ValueError(
                    f"Duplicate function name: {f_name}. Use unique function names, or use ToolsProvider and assign a unique alias to each function."
                )
        self.tool_name_map[f_name] = tool

    @deprecated("Use add_tool instead. This method will be removed in future versions.")
    def add_function(self, function: Union[Callable, BaseTool], alias: str = None):
        """Add function to ToolsProvider. If alias is provided, it will be used as function name in LLM"""
        self.add_tool(function, alias)

    def __contains__(self, function):
        return function in self.tools

    def get_function_schemas(self, inputs, _index: int = None):
        if self.tool_schemas:
            _f_schemas = []
            for i, (alias, f_schema_builder) in enumerate(
                zip(self.aliases, self.tool_schemas)
            ):
                if _index is not None and i != _index:
                    continue

                if callable(f_schema_builder) and not isinstance(
                    f_schema_builder, BaseTool
                ):
                    _f_schema = f_schema_builder(inputs)
                else:
                    _f_schema = f_schema_builder

                if alias:
                    _f_schema["name"] = alias

                _f_schemas.append(_f_schema)

            return _f_schemas
        else:
            None

    def get_function_schema(self, function: Union[str, Callable], inputs: dict):
        index = None
        if isinstance(function, str):
            func = self.tool_name_map[function]
        else:
            func = function

        _index = self.tools.index(func)
        return self.get_function_schemas(inputs, _index=_index)[0]

    def get_function(self, function_name: str = None):
        if function_name in self.tool_name_map:
            return self.tool_name_map[function_name]
        else:
            raise KeyError(f"Invalid function {function_name}")

    def __iter__(self):
        return iter(self.tools)

    def index(self, function):
        return self.tools.index(function)

    def get_tool_by_name(self, function_name: str, raise_errors: bool = True):
        """Get tool by name. If function_name is not provided, return the first tool."""

        if function_name in self.tool_name_map:
            return self.tool_name_map[function_name]
        elif raise_errors:
            raise KeyError(f"Invalid function {function_name}")


T = TypeVar("T")


class OutputWithFunctionCall(BaseModel, Generic[T]):
    output_text: str
    output_message: AIMessage
    output: T
    tool_call_id: Union[str, None] = None
    function_name: Union[str, None] = None
    function_arguments: Union[Dict[str, Any], str, None]
    function: Union[Callable, None] = None
    function_async: Union[Callable, None] = None
    result: Union[Any, None] = None
    _result_generated = PrivateAttr(False)

    @property
    def is_function_call(self):
        return bool(self.function or self.function_async)

    @property
    def support_async(self):
        return bool(self.function_async)

    @property
    def support_sync(self):
        return bool(self.function)

    async def execute_async(
        self, augment_args: Callable[[Dict[str, Any]], Dict[str, Any]] = None
    ):
        """Executes the function asynchronously."""
        if not (self.function or self.function_async):
            raise ValueError("No function to execute")
        if not (self.function or self.function_async):
            raise ValueError("No function to execute")

        if not self.function:
            raise ValueError("Function not set")
        else:
            if isinstance(self.function, BaseTool):
                if augment_args:
                    call_kwargs = augment_args(**self.function_arguments)
                else:
                    call_kwargs = self.function_arguments or {}
                result = await self.function.ainvoke(call_kwargs)
            else:
                call_args, call_kwargs = self._get_final_func_args(augment_args)
                result = self.function(*call_args, **call_kwargs)
                if asyncio.iscoroutine(result):
                    result = await result
        self.result = result
        self._result_generated = True
        return result

    def _get_final_func_args(self, augment_args: Callable) -> Dict[str, Any]:
        call_kwargs = self.function_arguments or {}
        call_args = []
        for k in list(call_kwargs):
            if k.startswith("__arg"):
                # skip private attributes
                call_args.append(call_kwargs.pop(k))
        if augment_args:
            call_kwargs = augment_args(**call_kwargs)
        return call_args, call_kwargs

    def execute(self, augment_args: Callable[[Dict[str, Any]], Dict[str, Any]] = None):
        """Executes the function synchronously.
        If the function is async, it will be executed in a event loop.
        """
        if not (self.function or self.function_async):
            raise ValueError("No function to execute")

        if not self.function:
            raise ValueError("Function not set")
        else:
            if isinstance(self.function, BaseTool):
                if augment_args:
                    call_kwargs = augment_args(**self.function_arguments)
                else:
                    call_kwargs = self.function_arguments or {}
                result = self.function.invoke(call_kwargs)
            else:
                call_args, call_kwargs = self._get_final_func_args(augment_args)
                result = self.function(*call_args, **call_kwargs)

        if asyncio.iscoroutine(result):
            try:
                current_loop = asyncio.get_running_loop()
            except RuntimeError:
                current_loop = None
            if current_loop:
                raise RuntimeError(
                    "Cannot execute async function synchronously. Please use execute_async() instead.",
                )
            else:
                logging.warning(
                    "Executing async function synchronously. This is not recommended. Consider using execute_async() instead."
                )
                result = asyncio.run(self.function_async(**self.function_arguments))
        self.result = result
        self._result_generated = True
        return result

    @property
    def function_call_message(self):
        """Returns the function call message"""
        if not self.is_function_call:
            raise ValueError(
                "Output was not a function call. You can test this with is_function_call property"
            )
        if self.output_message:
            return self.output_message

    def to_function_message(self, result=None):
        """
        Deprecated: Use function_output_to_message instead
        """
        logging.warning(
            "to_function_message is deprecated, use function_output_to_message instead"
        )
        return self.function_output_to_message(function_output=result)

    def function_output_to_message(self, function_output=None):
        """
        Converts the result of the functional call to a FunctionMessage...
        you can override the result collected via execute with your own by providing function_output

        Args:
            function_output (Any, optional): function output. If None, it the result collected via execute() or execute_async() will be used. (One of them must be called before).
        """
        if not function_output:
            if not self._result_generated:
                try:
                    self.result = self.execute()
                except RuntimeError as e:
                    if "Cannot execute async function synchronously." in str(e):
                        raise RuntimeError(
                            "Cannot execute async function synchronously. Please use await execute_async() to generate the output of the function first"
                        ) from e
            function_output = self.result
        if isinstance(function_output, BaseModel):
            function_output = function_output.json()
        elif not isinstance(function_output, str):
            function_output = repr(function_output)

        return ToolMessage(
            tool_call_id=self.tool_call_id,
            name=self.function_name,
            content=function_output,
        )


class ToolCall(BaseModel):
    """A tool call that can be used in a LLM chat session."""

    id: str
    name: str
    args: Union[Dict[str, Any], str, None] = None
    function: Optional[Callable] = None
    result: str = None
    is_error_result: bool = False
    metadata: Optional[Dict[str, Any]] = None
    _result_original_value: Any = PrivateAttr(None)

    def __init__(self, **data):
        if "args" in data and isinstance(data["args"], str):
            try:
                data["args"] = json.loads(data["args"])
            except json.JSONDecodeError:
                # If args is a string but not valid JSON, we leave it as is
                data["metadata"] = data.get("metadata", {})
                data["metadata"]["args_string"] = data["args"]
                data["args"] = None
        super().__init__(**data)

    def invoke(self, **kwargs):
        """Executes the tool call and sets the result + adds it to LlmChatSession if available"""
        _res = self(**kwargs)
        if asyncio.iscoroutine(_res):
            raise RuntimeError(
                "Cannot execute async function synchronously. Please use ainvoke() instead."
            )

        self.set_result(_res)
        return _res

    async def ainvoke(self, **kwargs):
        """
        Executes the tool call and sets the result ... regardless of whether the function is async or sync.
        """
        _res = await self(**kwargs)
        if asyncio.iscoroutine(_res):
            self._result_original_value = await _res
        else:
            self._result_original_value = _res

        self.set_result(_res)
        return _res

    def __call__(self, **kwargs):
        """Calls the function with the provided arguments."""
        if not self.function:
            raise ValueError("No function to execute")
        if isinstance(self.function, BaseTool):
            # Merge stored args with any explicit kwargs provided at call time
            call_kwargs = {}
            if isinstance(self.args, dict):
                call_kwargs.update(self.args)
            if kwargs:
                call_kwargs.update(kwargs)
            return self.function.invoke(call_kwargs)
        else:
            args = []
            for key in kwargs:
                if key.startswith("__arg"):
                    args.append(kwargs.pop(key))

            call_kwargs = self.args or {}
            call_kwargs.update(kwargs)
            return self.function(*args, **call_kwargs)

    def _serialize_result(
        self, tool_results: Union[BaseModel, dict, str, Exception, list, dict]
    ) -> str:
        """Serialize tool results to a JSON string."""
        if isinstance(tool_results, BaseModel):
            return tool_results.model_dump_json()
        elif isinstance(tool_results, dict):
            return json.dumps(tool_results)
        elif isinstance(tool_results, list):
            return "\n".join([self._serialize_result(item) for item in tool_results])
        elif isinstance(tool_results, Exception):
            return f"Error {tool_results.__class__.__name__}: {str(tool_results)}"
        else:
            return str(tool_results)

    def set_result(self, result: Union[BaseModel, dict, str, list]):
        if asyncio.iscoroutine(result):
            raise RuntimeError(
                "Result must be evaluated before setting it! Trying to set coroutine result"
            )
        self._result_original_value = result
        self.result = self._serialize_result(result)
        if LlmChatSession.get_current_session():
            LlmChatSession.get_current_session().add_message(self.to_tool_message())
        return self.result

    def set_error_result(self, error: Union[Exception, str]):
        self.is_error_result = True
        if isinstance(error, Exception):
            self.result = f"Error {error.__class__.__name__}: {str(error)}"
        else:
            self.result = str(error)
        if LlmChatSession.get_current_session():
            LlmChatSession.get_current_session().add_message(self.to_tool_message())

    def to_tool_message(
        self, result: Union[BaseModel, dict, str, list] = None
    ) -> ToolMessage:
        """Converts the tool call to a ToolMessage."""
        result = result or self.result

        return ToolMessage(
            tool_call_id=self.id, content=result or self.result, name=self.name
        )
