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

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

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 = True,
        **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,  # FIX: Cohérence avec complete()
        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.
        """
        ...

    @staticmethod
    def _ensure_messages_list(
        i: LanguageModelInput,
    ) -> List[Union[HumanMessage, AIMessage, ToolMessage]]:
        """Ensures that the input is always a list of messages."""
        if isinstance(i, str):
            return [HumanMessage(content=i)]
        return i

    def _execute_and_build_tool_messages(
        self, tool_calls: List[ToolCall]
    ) -> List[ToolMessage]:
        """
        Executes a list of ToolCall and returns a list of ToolMessage with the results.
        """
        tool_messages = []
        if not tool_calls:
            return tool_messages

        for tool_call in tool_calls:
            tool_name = tool_call.function.name
            tool_arguments_str = tool_call.function.arguments
            tool_call_id = tool_call.id

            try:
                tool_arguments = (
                    json.loads(tool_arguments_str) if tool_arguments_str else {}
                )
                tool_result = self._parent.tools.execute(tool_name, tool_arguments)

                result_content = (
                    json.dumps(tool_result)
                    if not isinstance(tool_result, str)
                    else tool_result
                )
                tool_messages.append(
                    ToolMessage(content=result_content, tool_call_id=tool_call_id)
                )

            except json.JSONDecodeError:
                error_msg = (
                    f"Error: Invalid JSON arguments provided for tool {tool_name}."
                )
                logger.error(f"{error_msg} Arguments: {tool_arguments_str}")
                tool_messages.append(
                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
                )
            except ToolExecutionError as e:
                error_msg = f"Error: {str(e)}"
                logger.error(f"Error executing tool {tool_name}: {error_msg}")
                tool_messages.append(
                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
                )
            except Exception as e:
                error_msg = f"Error: Unexpected error executing tool {tool_name}."
                logger.error(f"{error_msg}: {str(e)}", exc_info=True)
                tool_messages.append(
                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
                )

        return tool_messages

    @staticmethod
    def _process_stream_chunk(
        chunk: ChatCompletionChunk, reconstructed_tool_calls, assistant_content
    ):
        """
        Processes a chunk of stream and updates the data structures.
        Returns the updated assistant content and True if tool_calls are complete.
        """
        if not chunk.choices:
            return assistant_content, False

        delta: Delta = chunk.choices[0].delta
        if not delta:
            return assistant_content, False

        # Gather content
        if hasattr(delta, "reasoning_content") and delta.reasoning_content:
            assistant_content += delta.reasoning_content
        if delta.content:
            assistant_content += delta.content

        # Traiter les tool_calls
        if delta.tool_calls:
            for tc in delta.tool_calls:
                idx = tc.index
                # Étendre la liste si nécessaire
                while len(reconstructed_tool_calls) <= idx:
                    reconstructed_tool_calls.append(
                        {
                            "id": "",
                            "type": "function",
                            "function": {"name": "", "arguments": ""},
                        }
                    )

                target = reconstructed_tool_calls[idx]
                if tc.id:
                    target["id"] = tc.id
                if tc.function:
                    if tc.function.name:
                        target["function"]["name"] = tc.function.name
                    if tc.function.arguments:
                        target["function"]["arguments"] += tc.function.arguments

        # Vérifier si les tool_calls sont terminés
        tool_calls_finished = chunk.choices[0].finish_reason == "tool_calls"
        return assistant_content, tool_calls_finished

    def _reconstruct_tool_calls_from_stream(self, stream_response):
        """
        Common method for reconstructing tool_calls from a stream.
        Yields the chunks and returns the contents of the assistant and the reconstructed tool_calls.
        """
        reconstructed_tool_calls: List[dict] = []
        assistant_content = ""

        for chunk in stream_response:
            yield chunk  # Yield chunk for streaming

            assistant_content, tool_calls_finished = self._process_stream_chunk(
                chunk, reconstructed_tool_calls, assistant_content
            )

            # if tool_calls_finished yield final result
            if tool_calls_finished:
                tool_calls = [
                    ToolCall.model_validate(tc) for tc in reconstructed_tool_calls
                ]
                yield "tool_calls_result", assistant_content, tool_calls
                return

        yield "no_tool_calls", assistant_content, []

    async def _async_reconstruct_tool_calls_from_stream(self, stream_response):
        """
        Asynchronous version of reconstructing tool_calls from a stream.
        Yields chunks and reports results via special tuples.
        """
        reconstructed_tool_calls: List[dict] = []
        assistant_content = ""

        async for chunk in stream_response:
            yield chunk

            assistant_content, tool_calls_finished = self._process_stream_chunk(
                chunk, reconstructed_tool_calls, assistant_content
            )

            if tool_calls_finished:
                tool_calls = [
                    ToolCall.model_validate(tc) for tc in reconstructed_tool_calls
                ]
                yield "tool_calls_result", assistant_content, tool_calls
                return

        yield "no_tool_calls", assistant_content, []


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 = True,
        **kwargs,
    ) -> ChatCompletionResponse:
        messages = self._ensure_messages_list(input)

        if not auto_execute_tools:
            return endpoints.chat(
                self._client,
                messages,
                model,
                temperature,
                tools_to_payload(tools),
                **kwargs,
            )

        max_iterations = kwargs.pop("max_tool_iterations", 10)
        iteration = 0

        while iteration < max_iterations:
            response = endpoints.chat(
                self._client,
                messages,
                model,
                temperature,
                tools=tools_to_payload(tools),
                **kwargs,
            )

            tool_calls = response.choices[0].message.tool_calls
            if not tool_calls:
                return response

            # Add the AI message to the messages list
            ai_message = AIMessage(
                content=response.choices[0].message.content, tool_calls=tool_calls
            )
            messages.append(ai_message)

            # Execute tools and add result
            tool_messages = self._execute_and_build_tool_messages(tool_calls)
            messages.extend(tool_messages)

            iteration += 1

        logger.warning(f"Max tool iterations ({max_iterations}) reached")
        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]:
        messages = self._ensure_messages_list(input)

        if not auto_execute_tools:
            tool_payload = tools_to_payload(tools) if tools else None
            yield from endpoints.stream(
                self._client, messages, model, temperature, tools=tool_payload, **kwargs
            )
            return

        yield from self._run_with_tools_loop(
            messages=messages,
            model=model,
            temperature=temperature,
            tools=tools,
            **kwargs,
        )

    def _run_with_tools_loop(
        self,
        messages: List[Union[HumanMessage, AIMessage, ToolMessage]],
        model: str,
        temperature: float,
        tools: Optional[List[Tool]],
        **kwargs,
    ) -> Generator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None, None]:
        tool_payload = tools_to_payload(tools) if tools else None
        max_iterations = kwargs.pop("max_tool_iterations", 10)
        iteration = 0

        while iteration < max_iterations:
            # Stream response and build tool calls
            stream_response = endpoints.stream(
                self._client,
                messages,
                model,
                temperature,
                tools=tool_payload,
                full_chunk=True,
                **kwargs,
            )

            assistant_content = ""
            tool_calls = []

            for item in self._reconstruct_tool_calls_from_stream(stream_response):
                if isinstance(item, tuple) and item[0] == "tool_calls_result":
                    _, assistant_content, tool_calls = item
                    break
                elif isinstance(item, tuple) and item[0] == "no_tool_calls":
                    _, assistant_content, tool_calls = item
                    break
                else:
                    yield item  # C'est un chunk normal

            if not tool_calls:
                # Pas de tool_calls, fin de la conversation
                break

            # Ajouter le message AI avec tool_calls
            ai_message = AIMessage(
                content=assistant_content or None, tool_calls=tool_calls
            )
            messages.append(ai_message)

            # Exécuter les outils
            tool_messages = self._execute_and_build_tool_messages(tool_calls)
            messages.extend(tool_messages)

            iteration += 1

        if iteration >= max_iterations:
            logger.warning(
                f"Max tool iterations ({max_iterations}) reached in streaming mode"
            )


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 = True,
        **kwargs,
    ) -> ChatCompletionResponse:
        messages = self._ensure_messages_list(input)

        if not auto_execute_tools:
            resp = await async_endpoints.chat(
                self._client,
                messages,
                model,
                temperature,
                tools=tools_to_payload(tools),
                **kwargs,
            )
            return ChatCompletionResponse.model_validate(resp)

        max_iterations = kwargs.pop("max_tool_iterations", 10)
        iteration = 0

        while iteration < max_iterations:
            resp = await async_endpoints.chat(
                self._client,
                messages,
                model,
                temperature,
                tools=tools_to_payload(tools),
                **kwargs,
            )
            response = ChatCompletionResponse.model_validate(resp)

            tool_calls = response.choices[0].message.tool_calls
            if not tool_calls:
                return response

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

            tool_messages = self._execute_and_build_tool_messages(tool_calls)
            messages.extend(tool_messages)

            iteration += 1

        logger.warning(f"Max tool iterations ({max_iterations}) reached")
        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]:
        messages = self._ensure_messages_list(input)

        if not auto_execute_tools:
            async for chunk in async_endpoints.chat_stream(
                self._client,
                messages,
                model,
                temperature,
                tools=tools_to_payload(tools) if tools else None,
                full_chunk=full_chunk,
                **kwargs,
            ):
                yield chunk
            return

        # FIX: Mode agent unifié
        async for val in self._run_with_tools_loop(
            messages=messages,
            model=model,
            temperature=temperature,
            tools=tools,
            **kwargs,
        ):
            yield val

    async def _run_with_tools_loop(
        self,
        messages: List[Union[HumanMessage, AIMessage, ToolMessage]],
        model: str,
        temperature: float,
        tools: Optional[List[Tool]],
        **kwargs,
    ) -> AsyncGenerator[Union[str, ChatCompletionChunk, ChatCompletionResponse], None]:
        tool_payload = tools_to_payload(tools) if tools else None
        max_iterations = kwargs.pop("max_tool_iterations", 10)
        iteration = 0

        while iteration < max_iterations:
            stream_response = async_endpoints.chat_stream(
                self._client,
                messages,
                model,
                temperature,
                tools=tool_payload,
                **kwargs,
            )

            assistant_content = ""
            tool_calls = []

            async for item in self._async_reconstruct_tool_calls_from_stream(
                stream_response
            ):
                if isinstance(item, tuple) and item[0] == "tool_calls_result":
                    _, assistant_content, tool_calls = item
                    break
                elif isinstance(item, tuple) and item[0] == "no_tool_calls":
                    _, assistant_content, tool_calls = item
                    break
                else:
                    yield item

            if not tool_calls:
                break

            ai_message = AIMessage(
                content=assistant_content or None, tool_calls=tool_calls
            )
            messages.append(ai_message)

            tool_messages = self._execute_and_build_tool_messages(tool_calls)
            messages.extend(tool_messages)

            iteration += 1

        if iteration >= max_iterations:
            logger.warning(
                f"Max tool iterations ({max_iterations}) reached in async streaming mode"
            )
