from typing import Any

from sinapsis_chatbots_base.helpers.llm_keys import LLMChatKeys
from sinapsis_chatbots_base.helpers.postprocess_text import postprocess_text
from sinapsis_core.data_containers.data_packet import DataContainer, TextPacket

from sinapsis_vllm.templates.vllm_text_completion import vLLMTextCompletion


class vLLMBatchTextCompletion(vLLMTextCompletion):
    """Template for batched text completion using vLLM's continuous batching engine.

    Usage example:

    agent:
      name: batch_extraction_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: vLLMBatchTextCompletion
      class_name: vLLMBatchTextCompletion
      template_input: InputTemplate
      attributes:
        init_args:
          llm_model_name: 'unsloth/Qwen3-0.6B'
          max_model_len: 2048
          dtype: auto
          seed: 0
          gpu_memory_utilization: 0.9
          max_num_seqs: 256
          disable_log_stats: true
        completion_args:
          temperature: 0.2
          top_p: 0.8
          top_k: 20
          min_p: 0
          max_tokens: 1024
        system_prompt: You are a helpful AI assistant.
    """

    def _build_conversation(self, packet: TextPacket) -> tuple[list[dict], str, str | None]:
        """Builds a complete conversation message list for a single packet.

        Assembles the system prompt, chat history, RAG context, and user message
        into an ordered message list suitable for ``llm.chat()``.

        Args:
            packet (TextPacket): The incoming TextPacket to build the conversation for.

        Returns:
            tuple[list[dict], str, str | None]: A tuple of (message_list, user_id, session_id).
        """
        full_context: list[dict] = []
        user_id, session_id, prompt = self.prepare_conversation_context(packet)

        if self.system_prompt:
            full_context.append(self.generate_dict_msg(LLMChatKeys.system_value, self.system_prompt))

        if self.attributes.chat_history_key:
            full_context.extend(packet.generic_data.get(self.attributes.chat_history_key, []))

        full_context.append(self.generate_dict_msg(LLMChatKeys.user_value, prompt))
        return full_context, user_id, session_id

    def generate_response(self, container: DataContainer) -> DataContainer:
        """Processes all text packets in a single batched call to vLLM.

        Builds conversation message lists for each packet independently (preserving
        per-packet chat history and RAG context), then submits them all to
        ``llm.chat()`` in one batched call. Responses are zipped back to their
        corresponding packets with structured output parsing applied when configured.

        Args:
            container (DataContainer): Container with incoming messages.

        Returns:
            DataContainer: Updated DataContainer with the LLM responses.
        """
        self.logger.debug("Batched chatbot in progress")

        all_conversations: list[list[dict]] = []
        packet_metadata: list[tuple[str, str | None]] = []

        for packet in container.texts:
            conversation, user_id, session_id = self._build_conversation(packet)
            all_conversations.append(conversation)
            packet_metadata.append((user_id, session_id))

        if not all_conversations:
            return container

        self.logger.debug(f"Sending batch of {len(all_conversations)} conversations to vLLM")
        outputs = self.llm.chat(
            messages=all_conversations,
            sampling_params=self.sampling_params,
            use_tqdm=False,
        )

        responses = []
        for output, (user_id, session_id) in zip(outputs, packet_metadata):
            response = output.outputs[0].text if output.outputs else None

            if response:
                response = postprocess_text(response, self.attributes.pattern, self.attributes.keep_before)
                generic_data: dict[str, Any] = {}

                if self.attributes.completion_args.response_format.type_ == "json_object":
                    parsed = self._parse_json_response(response)
                    if parsed is not None:
                        generic_data[self.attributes.structured_output_key] = parsed

                responses.append(TextPacket(source=session_id, content=response, id=user_id, generic_data=generic_data))

        self.logger.debug(f"Batch complete: {len(responses)}/{len(all_conversations)} responses generated")
        container.texts.extend(responses)
        return container
