from abc import ABC, abstractmethod
from typing import Optional, Union, Generator, AsyncGenerator, List, Awaitable

from ..models.chat import (
    LanguageModelInput,
    HumanMessage,
    ChatCompletionResponse,
    AIMessage,
    ToolMessage,
    ChatCompletionChunk,
)
from ..models.tools import Tool, tools_to_payload
from ..exceptions import ToolExecutionError
import json
from ..api import async_endpoints, endpoints
from ..abstract.modules import AsyncModule, SyncModule
import logging

logger = logging.getLogger(__name__)


class AbstractChatModule(ABC):
    """
    Abstract base class for chat modules.
    Handles common logic for processing tool calls in conversational LLM flows.
    """

    def __init__(self, parent, http_client):
        """
        Initialize the AbstractChatModule.

        Args:
            parent: Parent client instance.
            http_client: HTTP client for making requests.
        """
        self._parent = parent
        self._client = http_client

    @abstractmethod
    def complete(
        self,
        input: LanguageModelInput,
        model: str,
        temperature: float = 0.7,
        tools: Optional[List[Tool]] = None,
        auto_execute_tools: bool = False,
        **kwargs,
    ) -> Union[ChatCompletionResponse, "Awaitable[ChatCompletionResponse]"]:
        """
        Sends a chat message (or list of messages) to the model and returns a structured response.
        Supports tool usage and optional automatic tool execution.

        Args:
            input (LanguageModelInput): Message(s) to send to the model
            model (str): Identifier of the language model to use.
            temperature (float): Sampling temperature for response creativity (0.0–1.0, default: 0.7).
            tools (Optional[list]): List of tools that can be invoked by the model.
            auto_execute_tools (bool): If True, automatically executes tool calls generated by the model (default: False).
            **kwargs: Additional arguments passed to the API.

        Returns:
            ChatCompletionResponse: Structured model response.
        """
        ...

    @abstractmethod
    def stream(
        self,
        input: LanguageModelInput,
        model: str,
        temperature: float = 0.7,
        tools: Optional[List[Tool]] = None,
        auto_execute_tools: bool = True,
        full_chunk: bool = True,
        **kwargs,
    ) -> Union[
        Generator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None, None],
        AsyncGenerator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None],
    ]:
        """
        Streams a chat completion response from the model. If tools are provided,
        makes an initial non-streaming call to determine tool invocations,
        executes them if auto_execute_tools is enabled, and then streams the final response.

        Args:
            input (LanguageModelInput): Message(s) to send to the model
            model (str): Identifier of the language model to use.
            temperature (float): Sampling temperature for response creativity (0.0–1.0, default: 0.7).
            tools (Optional[list]): List of tools that can be invoked by the model.
            auto_execute_tools (bool): If True, automatically executes tool calls generated by the model (default: False).
            full_chunk (bool): If True, yields full response objects; if False, yields only the content (default: True).
            **kwargs: Additional keyword arguments for the API call.

        Yields:
            Union[str, ChatCompletionChunk, ChatCompletionResponse]: Response chunks or final responses as they become available.
        """
        ...

    @abstractmethod
    def _handle_tool_calls(
        self,
        response: ChatCompletionResponse,
        input,
        model,
        temperature,
        stream: bool = False,
        full_chunk: bool = True,
        **kwargs,
    ):
        """
        Handles execution of tool calls found in a model response. Automatically executes the requested tools and
        integrates their results into the conversation, then continues the chat either in streaming or blocking mode.

        Args:
            response (ChatCompletionResponse): The initial model response containing tool calls.
            input: Original input messages.
            model: Model identifier.
            temperature: Sampling temperature.
            stream: If True, returns a streaming generator.
            full_chunk: (Streaming only) If True, yields full chunks.
            **kwargs: Additional arguments for the next API call.

        Returns:
            ChatCompletionResponse or AsyncGenerator: If stream is True, returns a generator yielding tool results and final response.
            Otherwise, returns a complete ChatCompletionResponse with tool results integrated.
        """
        ...

    def _build_tool_messages(
        self,
        response: ChatCompletionResponse,
        input: Union[LanguageModelInput, List[HumanMessage]],
    ):
        """
        Utility: Build the message list including tool calls and their results.
        """
        tool_calls = response.choices[0].message.tool_calls
        messages: list = (
            input if isinstance(input, list) else [HumanMessage(content=input)]
        )

        messages.append(
            AIMessage(
                content=response.choices[0].message.content or "",
                tool_calls=tool_calls,
            )
        )

        for tool_call in tool_calls or []:
            tool_name = tool_call.function.name
            tool_arguments_str = tool_call.function.arguments
            try:
                tool_arguments = (
                    json.loads(tool_arguments_str) if tool_arguments_str else {}
                )
                tool_result = self._parent.tools.execute(tool_name, tool_arguments)
                messages.append(
                    ToolMessage(
                        content=json.dumps(tool_result)
                        if not isinstance(tool_result, str)
                        else tool_result,
                        tool_call_id=tool_call.id,
                    )
                )
            except json.JSONDecodeError:
                logger.error(
                    f"Invalid JSON arguments provided for tool {tool_name}."
                    f" Arguments: {tool_arguments_str}"
                )
                messages.append(
                    ToolMessage(
                        content=f"Error: Invalid JSON arguments provided for tool {tool_name}.",
                        tool_call_id=tool_call.id,
                    )
                )
            except ToolExecutionError as e:
                logger.error(f"Error executing tool {tool_name}: {str(e)}")
                messages.append(
                    ToolMessage(
                        content=f"Error: {str(e)}",
                        tool_call_id=tool_call.id,
                    )
                )
            except Exception as e:
                logger.error(f"Unexpected error executing tool {tool_name}: {str(e)}")
                messages.append(
                    ToolMessage(
                        content=f"Error: Unexpected error executing tool {tool_name}.",
                        tool_call_id=tool_call.id,
                    )
                )

        return messages


class SyncChatModule(AbstractChatModule, SyncModule):
    """
    `ChatModule` is a synchronous module that provides a high-level interface for interacting with
    language models in a conversational format. It supports standard completions, streaming responses,
    and automatic tool execution within conversations. The module abstracts all communication with the backend API,
    providing both blocking and generator-based usage.
    """

    def __init__(self, parent, http_client):
        super().__init__(parent, http_client)
        SyncModule.__init__(self, parent, http_client)

    def complete(
        self,
        input: LanguageModelInput,
        model: str,
        temperature: float = 0.7,
        tools: Optional[list[Tool]] = None,
        auto_execute_tools: bool = False,
        **kwargs,
    ) -> ChatCompletionResponse:
        response = endpoints.chat(
            self._client, input, model, temperature, tools_to_payload(tools), **kwargs
        )

        if (
            auto_execute_tools
            and response.choices
            and response.choices[0].message.tool_calls
        ):
            return self._handle_tool_calls(
                input=input,
                response=response,
                model=model,
                temperature=temperature,
                stream=False,
                **kwargs,
            )

        return response

    def stream(
        self,
        input: LanguageModelInput,
        model: str,
        temperature: float = 0.7,
        tools: Optional[list[Tool]] = None,
        auto_execute_tools: bool = True,
        full_chunk: bool = True,
        **kwargs,
    ) -> Generator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None, None]:
        if tools:
            resp = endpoints.chat(
                self._client,
                input,
                model,
                temperature,
                tools_to_payload(tools),
                **kwargs,
            )
            response = ChatCompletionResponse.model_validate(resp)

            if not auto_execute_tools:
                yield response if full_chunk else response.choices[0].message.content
                return

            yield from self._handle_tool_calls(
                response,
                input,
                model,
                temperature,
                stream=True,
                full_chunk=full_chunk,
                **kwargs,
            )
        else:
            for chunk in endpoints.stream(
                self._client,
                input,
                model,
                temperature,
                full_chunk=full_chunk,
                **kwargs,
            ):
                yield ChatCompletionChunk.model_validate(chunk) if full_chunk else chunk

    def _handle_tool_calls(
        self,
        response: ChatCompletionResponse,
        input,
        model,
        temperature,
        stream: bool = False,
        full_chunk: bool = True,
        **kwargs,
    ):
        tool_calls = response.choices[0].message.tool_calls

        if tool_calls is None:
            return response

        messages = self._build_tool_messages(response, input)

        if stream:
            return self.stream(
                messages, model, temperature, full_chunk=full_chunk, **kwargs
            )
        else:
            return self.complete(messages, model, temperature, **kwargs)


class AsyncChatModule(AbstractChatModule, AsyncModule):
    """
    `ChatModule` is an asynchronous module that provides a high-level interface for interacting with
    language models in a conversational format. It supports standard completions, streaming responses,
    and automatic tool execution within conversations. The module abstracts all communication with the backend API,
    providing both blocking and generator-based usage.
    """

    def __init__(self, parent, http_client):
        super().__init__(parent, http_client)
        AsyncModule.__init__(self, parent, http_client)

    async def complete(
        self,
        input: LanguageModelInput,
        model: str,
        temperature: float = 0.7,
        tools: Optional[list[Tool]] = None,
        auto_execute_tools: bool = False,
        **kwargs,
    ) -> ChatCompletionResponse:
        resp = await async_endpoints.chat(
            self._client, input, model, temperature, tools_to_payload(tools), **kwargs
        )
        response = ChatCompletionResponse.model_validate(resp)

        if (
            auto_execute_tools
            and response.choices
            and response.choices[0].message.tool_calls
        ):
            return await self._handle_tool_calls(
                input=input,
                response=response,
                model=model,
                temperature=temperature,
                stream=False,
                **kwargs,
            )

        return response

    async def stream(
        self,
        input: LanguageModelInput,
        model: str,
        temperature: float = 0.7,
        tools: Optional[list[Tool]] = None,
        auto_execute_tools: bool = True,
        full_chunk: bool = True,
        **kwargs,
    ) -> AsyncGenerator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None]:
        if tools:
            resp = await async_endpoints.chat(
                self._client,
                input,
                model,
                temperature,
                tools_to_payload(tools),
                **kwargs,
            )
            response = ChatCompletionResponse.model_validate(resp)

            if not auto_execute_tools:
                if full_chunk:
                    yield response
                else:
                    yield response.choices[0].message.content
                return

            async for val in await self._handle_tool_calls(
                response,
                input,
                model,
                temperature,
                stream=True,
                full_chunk=full_chunk,
                **kwargs,
            ):
                yield val
        else:
            async for chunk in async_endpoints.chat_stream(
                self._client,
                input,
                model,
                temperature,
                full_chunk=full_chunk,
                **kwargs,
            ):
                yield ChatCompletionChunk.model_validate(chunk) if full_chunk else chunk

    async def _handle_tool_calls(
        self,
        response: ChatCompletionResponse,
        input,
        model,
        temperature,
        stream: bool = False,
        full_chunk: bool = True,
        **kwargs,
    ):
        tool_calls = response.choices[0].message.tool_calls

        if tool_calls is None:
            return response

        messages = self._build_tool_messages(response, input)

        if stream:
            return self.stream(
                messages, model, temperature, full_chunk=full_chunk, **kwargs
            )
        else:
            return await self.complete(messages, model, temperature, **kwargs)
