"""Module that defines the ResponseWrapper class."""

import logging
from typing import Iterator, Optional

from openai.types.chat import ChatCompletion, ChatCompletionMessage

from encourage.llm.response import Response
from encourage.llm.vllm_classes import RequestOutput
from encourage.prompts.context import Context
from encourage.prompts.conversation import Conversation, Role
from encourage.prompts.meta_data import MetaData
from encourage.prompts.prompt import Prompt
from encourage.prompts.prompt_collection import PromptCollection
from encourage.utils.tracing import enable_tracing

# Configure logging
logger = logging.getLogger(__name__)


class ResponseWrapper:
    """Class that aggregates ChatCompletion with corresponding PromptCollection details."""

    @classmethod
    def from_prompt_collection(
        cls,
        chat_completions: list[ChatCompletion],
        collection: PromptCollection,
    ) -> "ResponseWrapper":
        """Create ResponseWrapper from ChatCompletion and PromptCollection or Conversation."""
        if len(chat_completions) != len(collection.prompts):
            raise ValueError("The number of request outputs does not match the number of prompts.")

        # Create responses for each chat_completion and its corresponding prompt
        responses = [
            cls.handle_prompt_response(chat_completion, prompt)
            for chat_completion, prompt in zip(chat_completions, collection.prompts)
        ]
        return cls(responses)

    @classmethod
    @enable_tracing(span_name="handle_prompt_response")
    def handle_prompt_response(cls, chat_completion: ChatCompletion, prompt: Prompt) -> Response:
        """Create a Response object from a ChatCompletion and Prompt."""
        prompt.conversation = transform_chat_completion_to_conversation(prompt.conversation)

        return Response(
            request_id=chat_completion.id,
            prompt_id=str(prompt.id),
            conversation_id=0,
            sys_prompt=prompt.conversation.sys_prompt,
            user_prompt=prompt.conversation.get_last_message_by_user(),
            response=chat_completion.choices[0].message.content
            if chat_completion.choices
            else "No response",
            meta_data=prompt.meta_data,
            context=prompt.context,
            arrival_time=(chat_completion.created if chat_completion.created is not None else 0.0),
            finished_time=0.0,
        )

    @classmethod
    def from_conversation(
        cls,
        chat_completions: list[ChatCompletion],
        conversation: Conversation,
        contexts: list[Context] = [],
        meta_datas: list[MetaData] = [],
    ) -> "ResponseWrapper":
        """Create ResponseWrapper from ChatCompletion and Conversation."""
        # Check for mismatched lengths
        if len(chat_completions) != len(conversation.get_messages_by_role(Role.USER)):
            raise ValueError("The number of request outputs does not match the number of messages.")

        # Safely access the system prompt, handle missing system messages
        sys_messages = conversation.get_messages_by_role(Role.SYSTEM)
        if not sys_messages:
            raise ValueError("No system messages found in the conversation.")
        sys_prompt = sys_messages[0]["content"]

        # Extract user messages
        user_messages = [msg["content"] for msg in conversation.get_messages_by_role(Role.USER)]

        responses = []
        meta_datas = meta_datas or [MetaData()] * len(chat_completions)
        contexts = contexts or [Context()] * len(chat_completions)
        for conversation_id, (chat_completion, user_message, meta_data, context) in enumerate(
            zip(chat_completions, user_messages, meta_datas, contexts)
        ):
            response = cls.handle_conversation_response(
                sys_prompt, conversation_id, chat_completion, user_message, context, meta_data
            )
            responses.append(response)

        return cls(responses)

    @staticmethod
    @enable_tracing(span_name="handle_conversation_response")
    def handle_conversation_response(
        sys_prompt: str,
        conversation_id: int,
        chat_completion: ChatCompletion,
        message: str,
        context: Context,
        meta_data: MetaData,
    ) -> Response:
        """Create a Response object from a ChatCompletion and Conversation."""
        return Response(
            request_id=chat_completion.id,
            prompt_id="",
            conversation_id=conversation_id,
            sys_prompt=sys_prompt,
            user_prompt=message,
            response=chat_completion.choices[0].message.content
            if chat_completion.choices
            else "No response",
            context=context,
            meta_data=meta_data,
            arrival_time=(chat_completion.created if chat_completion.created is not None else 0.0),
            finished_time=0.0,
        )

    @classmethod
    def from_request_output(
        cls,
        request_outputs: list[RequestOutput],
        collection: PromptCollection,
    ) -> "ResponseWrapper":
        """Create ResponseWrapper from RequestOutput and PromptCollection."""
        if len(request_outputs) != len(collection.prompts):
            raise ValueError("The number of request outputs does not match the number of prompts.")

        # Create responses for each chat_completion and its corresponding prompt
        responses = [
            cls.handle_request_outputs(request_output, prompt)
            for request_output, prompt in zip(request_outputs, collection.prompts)
        ]
        return cls(responses)

    @staticmethod
    def handle_request_outputs(
        request_output: RequestOutput,
        prompt: Prompt,
    ) -> Response:
        """Handle a single request output and its corresponding prompt."""
        # Store a list of output texts in the response
        if not request_output.outputs:
            response_texts = "No response"
        elif len(request_output.outputs) == 1:
            response_texts = request_output.outputs[0].text
        else:
            response_texts = [o.text for o in request_output.outputs]

        return Response(
            request_id=request_output.request_id,
            prompt_id=str(prompt.id),
            conversation_id=0,
            sys_prompt=prompt.conversation.sys_prompt,
            user_prompt=prompt.conversation.get_last_message_by_user(),
            response=response_texts,
            meta_data=prompt.meta_data,
            context=prompt.context,
            arrival_time=0.0,
            finished_time=0.0,
        )

    def get_responses(self) -> list[str]:
        """Return the list of actually text responses from the model."""
        return [response.response for response in self.response_data]

    def print_response_summary(self) -> None:
        """Prints a summary of all responses."""
        for response in self.response_data:
            response.print_response()

    def print_response_by_prompt_id(self, prompt_id: str) -> None:
        """Print the response details for a specific prompt ID."""
        response = self._find_response_by("prompt_id", prompt_id)
        if response:
            response.print_response()
        else:
            print(f"No response found for Prompt ID: {prompt_id}")

    def print_response_by_request_id(self, request_id: str) -> None:
        """Print the response details for a specific request ID."""
        response = self._find_response_by("request_id", request_id)
        if response:
            response.print_response()
        else:
            logger.error(f"No response found for Request ID: {request_id}")

    def get_response_by_prompt_id(self, prompt_id: str) -> Optional[Response]:
        """Return the response details for a specific prompt ID."""
        return self._find_response_by("prompt_id", prompt_id)

    def get_response_by_request_id(self, request_id: str | int) -> Response | None:
        """Return the response details for a specific request ID."""
        return self._find_response_by("request_id", str(request_id))

    def _find_response_by(self, key: str, value: str) -> Response | None:
        """Helper method to find a response by a given key and value."""
        for response in self.response_data:
            if getattr(response, key) == value:
                return response
        return None

    def __repr__(self) -> str:
        return f"ResponseWrapper({self.response_data})"

    def __init__(self, responses: list[Response]):
        self.response_data = responses

    def __iter__(self) -> Iterator[Response]:
        """Allows iteration over the responses."""
        return iter(self.response_data)

    def __getitem__(self, key: int) -> Response:
        """Returns the response at the given index."""
        if not self.response_data:
            raise IndexError("Response data is empty.")
        if key < 0 or key >= len(self.response_data):
            raise IndexError("Index out of range.")
        return self.response_data[key]

    def __len__(self) -> int:
        """Returns the number of responses."""
        return len(self.response_data)


def transform_chat_completion_to_conversation(conversation: Conversation) -> Conversation:
    """Transforms a Conversation object to a ChatCompletion object."""
    # Create a new conversation object
    new_dialog = []
    for message in conversation.dialog:
        if isinstance(message, ChatCompletionMessage):  # type: ignore
            message = {"role": Role.TOOL, "content": message.tool_calls[0].function.arguments}  # type: ignore
        new_dialog.append(message)
    conversation.dialog = new_dialog
    return conversation
