from dataclasses import dataclass, field
import json
import logging
from typing import TYPE_CHECKING, List, Optional, Union

import openai
from pydantic import BaseModel

from rustic_ai.core.guild.agent import Agent, ProcessContext
from rustic_ai.core.guild.agent_ext.depends.llm.llm import LLM
from rustic_ai.core.guild.agent_ext.depends.llm.models import (
    ChatCompletionError,
    ChatCompletionRequest,
    ChatCompletionResponse,
    DiscriminatedLLMMessage,
    LLMMessage,
    ResponseCodes,
    SystemMessage,
)
from rustic_ai.llm_agent.llm_plugin_mixin import PluginConfigProtocol

if TYPE_CHECKING:
    from rustic_ai.llm_agent.llm_agent_conf import LLMAgentConfig


@dataclass
class LLMCompletionResult:
    """
    Result of calling the LLM with the plugin pipeline.

    This dataclass holds the response from the LLM along with any messages
    generated by postprocessors, without performing any sending operations.
    Agents are responsible for deciding how to send these results.
    """

    response: Union[ChatCompletionResponse, ChatCompletionError]
    """The response from the LLM or an error if the call failed."""

    final_request: ChatCompletionRequest
    """The final request sent to the LLM after preprocessing."""

    plugin_messages: List[BaseModel] = field(default_factory=list)
    """Messages generated by postprocessors."""

    @property
    def is_error(self) -> bool:
        """Check if the result is an error."""
        return isinstance(self.response, ChatCompletionError)

    @property
    def has_tool_calls(self) -> bool:
        """Check if the response contains tool calls."""
        if self.is_error:
            return False
        if not self.response.choices:
            return False
        return bool(self.response.choices[0].message.tool_calls)

    @property
    def reasoning(self) -> Optional[str]:
        """Extract reasoning content from the response if available."""
        if self.is_error:
            return None
        if not self.response.choices:
            return None
        return self.response.choices[0].message.reasoning_content


class LLMAgentHelper:
    """
    Helper class for LLM-based agents providing a plugin pipeline.

    This helper provides methods for preprocessing requests, calling the LLM,
    and postprocessing responses. It is designed to be used by both LLMAgent
    and ReActAgent with different sending behaviors.

    The plugin pipeline consists of:
    1. Request preprocessors - modify the request before the LLM call
    2. LLM call wrappers - preprocess before and postprocess after the LLM call
    3. Response postprocessors - act on the response after the LLM call

    Key methods:
    - preprocess_request(): Run preprocessing only, return prepared request
    - postprocess_response(): Run postprocessing only, return plugin messages
    - call_llm_with_plugins(): Full pipeline, return LLMCompletionResult (no sending)
    - invoke_llm_completion(): Backward-compatible method that sends plugin messages
    """

    @staticmethod
    def preprocess_request(
        agent: Agent,
        config: PluginConfigProtocol,
        llm: LLM,
        ctx: ProcessContext[ChatCompletionRequest],
        request: ChatCompletionRequest,
        llm_params: Optional[dict] = None,
    ) -> ChatCompletionRequest:
        """
        Run preprocessing on the request without calling the LLM.

        This method:
        1. Runs all request preprocessors in order
        2. Runs all LLM call wrapper preprocessors in order
        3. Merges config parameters with the request

        Args:
            agent: The agent instance
            config: Configuration with plugin lists
            llm: The LLM dependency (passed to plugins)
            ctx: The process context
            request: The initial ChatCompletionRequest
            llm_params: Optional dict of LLM parameters to merge (from config.get_llm_params())

        Returns:
            The preprocessed ChatCompletionRequest ready for the LLM
        """
        pre_processors = config.request_preprocessors or []
        wrap_processors = config.llm_request_wrappers or []

        # Run request preprocessors
        for pre_processor in pre_processors:
            request = pre_processor.preprocess(agent=agent, ctx=ctx, request=request, llm=llm)

        # Run wrapper preprocessors
        for wrap_processor in wrap_processors:
            request = wrap_processor.preprocess(agent=agent, ctx=ctx, request=request, llm=llm)

        # Merge config params with request
        if llm_params:
            request_dict = request.model_dump(exclude_none=True)
            final_dict = {**llm_params, **request_dict}
            request = ChatCompletionRequest.model_validate(final_dict)

        return request

    @staticmethod
    def postprocess_response(
        agent: Agent,
        config: PluginConfigProtocol,
        llm: LLM,
        ctx: ProcessContext[ChatCompletionRequest],
        final_request: ChatCompletionRequest,
        response: ChatCompletionResponse,
    ) -> List[BaseModel]:
        """
        Run postprocessing on the response without sending any messages.

        This method:
        1. Runs all LLM call wrapper postprocessors in reverse order
        2. Runs all response postprocessors in order
        3. Collects and returns all plugin-generated messages

        Args:
            agent: The agent instance
            config: Configuration with plugin lists
            llm: The LLM dependency (passed to plugins)
            ctx: The process context
            final_request: The request that was sent to the LLM
            response: The response from the LLM

        Returns:
            List of messages generated by postprocessors
        """
        wrap_processors = config.llm_request_wrappers or []
        post_processors = config.response_postprocessors or []

        plugin_messages: List[BaseModel] = []

        # Run wrapper postprocessors in reverse order
        for wrap_processor in reversed(wrap_processors):
            post_msg = wrap_processor.postprocess(
                agent=agent,
                ctx=ctx,
                final_prompt=final_request,
                llm_response=response,
                llm=llm,
            )
            if post_msg:
                plugin_messages.extend(post_msg)

        # Run response postprocessors in order
        for post_processor in post_processors:
            post_msg = post_processor.postprocess(
                agent=agent,
                ctx=ctx,
                final_prompt=final_request,
                llm_response=response,
                llm=llm,
            )
            if post_msg:
                plugin_messages.extend(post_msg)

        return plugin_messages

    @staticmethod
    def call_llm_with_plugins(
        agent: Agent,
        config: PluginConfigProtocol,
        llm: LLM,
        ctx: ProcessContext[ChatCompletionRequest],
        request: ChatCompletionRequest,
        llm_params: Optional[dict] = None,
        model: Optional[str] = None,
    ) -> LLMCompletionResult:
        """
        Execute the full plugin pipeline: preprocess, call LLM, postprocess.

        This method does NOT send any messages. The caller (agent) is responsible
        for deciding how to send the response and plugin messages.

        Args:
            agent: The agent instance
            config: Configuration with plugin lists and retry settings
            llm: The LLM dependency
            ctx: The process context
            request: The initial ChatCompletionRequest
            llm_params: Optional dict of LLM parameters to merge
            model: Optional model override

        Returns:
            LLMCompletionResult containing response, plugin messages, and final request
        """
        # Preprocess
        final_request = LLMAgentHelper.preprocess_request(
            agent=agent,
            config=config,
            llm=llm,
            ctx=ctx,
            request=request,
            llm_params=llm_params,
        )

        # Normalize and deduplicate messages before calling the LLM
        final_request = LLMAgentHelper._normalize_messages(final_request)

        # Call LLM
        response = llm.completion(final_request, model)

        # Postprocess with retry logic
        return LLMAgentHelper._postprocess_with_retry(
            agent=agent,
            config=config,
            llm=llm,
            ctx=ctx,
            final_request=final_request,
            response=response,
            model=model,
            max_retries=config.max_retries,
        )

    @staticmethod
    def _normalize_messages(request: ChatCompletionRequest) -> ChatCompletionRequest:
        """
        Normalize messages in the request.

        1. Extracts all system messages, deduplicates them, combines into a single
           system message, and prepends it to the list. This ensures models that
           require system messages at the beginning (e.g., Qwen) work correctly.
        2. Deduplicates non-system messages to save tokens when history enrichment
           causes overlapping messages.
        """
        system_contents: List[str] = []
        system_seen: set = set()
        other_messages: List[DiscriminatedLLMMessage] = []
        msg_seen: set = set()

        for msg in request.messages:
            if isinstance(msg, SystemMessage):
                content = msg.content
                if content not in system_seen:
                    system_seen.add(content)
                    system_contents.append(content)
            else:
                msg_key = json.dumps(msg.model_dump(exclude_none=True), sort_keys=True)
                if msg_key not in msg_seen:
                    msg_seen.add(msg_key)
                    other_messages.append(msg)

        if not system_contents and len(other_messages) == len(
            [m for m in request.messages if not isinstance(m, SystemMessage)]
        ):
            return request

        normalized_messages: List[DiscriminatedLLMMessage] = list(other_messages)
        if system_contents:
            combined_system = SystemMessage(content="\n".join(system_contents))
            normalized_messages.insert(0, combined_system)

        return request.model_copy(update={"messages": normalized_messages})

    @staticmethod
    def _postprocess_with_retry(
        agent: Agent,
        config: PluginConfigProtocol,
        llm: LLM,
        ctx: ProcessContext[ChatCompletionRequest],
        final_request: ChatCompletionRequest,
        response: ChatCompletionResponse,
        model: Optional[str],
        max_retries: int,
    ) -> LLMCompletionResult:
        """
        Run postprocessing with retry logic on failure.

        If postprocessing fails and retries are available, retries the LLM call.
        """
        try:
            plugin_messages = LLMAgentHelper.postprocess_response(
                agent=agent,
                config=config,
                llm=llm,
                ctx=ctx,
                final_request=final_request,
                response=response,
            )
            return LLMCompletionResult(
                response=response,
                final_request=final_request,
                plugin_messages=plugin_messages,
            )
        except Exception as e:  # pragma: no cover
            logging.error(f"Error in post processing: {e}")
            if max_retries > 0:
                logging.info(f"Retrying LLM call, remaining retries: {max_retries}")
                # Retry the LLM call
                response = llm.completion(final_request, model)
                return LLMAgentHelper._postprocess_with_retry(
                    agent=agent,
                    config=config,
                    llm=llm,
                    ctx=ctx,
                    final_request=final_request,
                    response=response,
                    model=model,
                    max_retries=max_retries - 1,
                )
            else:
                error = ChatCompletionError(
                    status_code=ResponseCodes.RESPONSE_PROCESSING_ERROR,
                    message=str(e),
                    model=agent.name,
                    request_messages=ctx.payload.messages if ctx.payload else [],
                )
                return LLMCompletionResult(
                    response=error,
                    final_request=final_request,
                    plugin_messages=[],
                )

    # =========================================================================
    # Backward-compatible methods (used by LLMAgent)
    # =========================================================================

    @staticmethod
    def invoke_llm_completion(
        agent: Agent,
        config: "LLMAgentConfig",
        llm: LLM,
        ctx: ProcessContext[ChatCompletionRequest],
        prompt: ChatCompletionRequest,
    ) -> Union[ChatCompletionResponse, ChatCompletionError]:
        """
        Invoke the LLM completion with the given context.

        This method maintains backward compatibility by calling the new
        plugin pipeline and sending plugin-generated messages.

        The fields from the chat completion request, Agent Config, and the LLM are combined.
        The LLM Configuration is used as the base, overwritten by Agent config and then
        the Chat Completion Request.

        Note: This method sends plugin messages but NOT the LLM response.
        Use call_llm_with_plugins() for full control over sending behavior.

        Args:
            agent: The agent instance
            config: LLMAgentConfig with plugin lists and LLM parameters
            llm: The LLM dependency
            ctx: The process context
            prompt: The ChatCompletionRequest to send

        Returns:
            ChatCompletionResponse or ChatCompletionError
        """
        result = LLMAgentHelper.call_llm_with_plugins(
            agent=agent,
            config=config,
            llm=llm,
            ctx=ctx,
            request=prompt,
            llm_params=config.get_llm_params(),
            model=config.model,
        )

        # Send plugin messages (backward-compatible behavior)
        if result.plugin_messages:
            reasoning = result.reasoning or ""
            for msg in result.plugin_messages:
                if reasoning:
                    ctx.send(msg, reason=reasoning)
                else:
                    ctx.send(msg)

        return result.response

    @staticmethod
    def process_api_status_error(
        model_name: str,
        ctx: ProcessContext[ChatCompletionRequest],
        status_code: ResponseCodes,
        error: openai.APIStatusError,
        messages: List[Union[LLMMessage]],
    ) -> ChatCompletionError:  # pragma: no cover
        """
        Process API status error and return a ChatCompletionError object.
        """
        cc_error = ChatCompletionError(
            status_code=status_code,
            message=error.message,
            response=error.response.text if error.response else None,
            model=model_name,
            request_messages=messages,
            body=error.body if hasattr(error, "body") else None,
        )

        return cc_error

    @staticmethod
    def invoke_llm_and_handle_response(
        agent: Agent,
        config: "LLMAgentConfig",
        llm: LLM,
        ctx: ProcessContext[ChatCompletionRequest],
        prompt: ChatCompletionRequest,
    ) -> Union[ChatCompletionResponse, ChatCompletionError]:
        """
        Invoke the LLM and handle the response with full sending behavior.

        This is a convenience method that:
        1. Calls the LLM with the plugin pipeline
        2. Sends plugin-generated messages
        3. Sends the LLM response (or error)

        This method is provided for backward compatibility. For more control,
        use call_llm_with_plugins() and handle sending in the agent.

        Args:
            agent: The agent instance
            config: LLMAgentConfig with plugin lists and LLM parameters
            llm: The LLM dependency
            ctx: The process context
            prompt: The ChatCompletionRequest to send

        Returns:
            ChatCompletionResponse or ChatCompletionError
        """
        try:
            result = LLMAgentHelper.call_llm_with_plugins(
                agent=agent,
                config=config,
                llm=llm,
                ctx=ctx,
                request=prompt,
                llm_params=config.get_llm_params(),
                model=config.model,
            )

            # Send plugin messages
            for msg in result.plugin_messages:
                if result.reasoning:
                    ctx.send(msg, reason=result.reasoning)
                else:
                    ctx.send(msg)

            # Send response or error
            if result.is_error:
                ctx.send_error(result.response)
            elif config.send_response:
                if result.reasoning:
                    ctx.send(result.response, reason=result.reasoning)
                else:
                    ctx.send(result.response)

            return result.response

        except openai.APIStatusError as e:  # pragma: no cover
            logging.error(f"Error in LLM completion: {e}")
            error = LLMAgentHelper.process_api_status_error(
                model_name=config.model,
                ctx=ctx,
                status_code=ResponseCodes(e.status_code),
                error=e,
                messages=ctx.payload.messages,
            )
            ctx.send_error(error)
            return error

        except Exception as e:  # pragma: no cover
            logging.error(f"Unexpected error in LLM completion: {e}")
            error = ChatCompletionError(
                status_code=ResponseCodes.INTERNAL_SERVER_ERROR,
                message=str(e),
                model=agent.name,
                request_messages=ctx.payload.messages,
            )
            ctx.send_error(error)
            return error
