import uuid
from collections.abc import AsyncGenerator
from copy import deepcopy

from sinapsis_chatbots_base.helpers.llm_keys import LLMChatKeys
from sinapsis_chatbots_base.helpers.postprocess_text import postprocess_text
from sinapsis_chatbots_base.helpers.tags import Tags
from sinapsis_core.data_containers.data_packet import DataContainer, TextPacket
from sinapsis_core.template_base.base_models import OutputTypes, UIPropertiesMetadata
from sinapsis_core.utils.sentinels import SENTINEL_GENERIC_KEY, SINAPSIS_END_OF_STREAM
from vllm import AsyncLLMEngine
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.sampling_params import RequestOutputKind

from sinapsis_vllm.templates.vllm_text_completion import vLLMTextCompletion


class vLLMStreamingTextCompletion(vLLMTextCompletion):
    """A streaming version of the vLLMTextCompletion template.

    Usage example:

    agent:
      name: my_streaming_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: vLLMStreamingTextCompletion
      class_name: vLLMStreamingTextCompletion
      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: 4
          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.
    """

    UIProperties = UIPropertiesMetadata(
        category="vLLM",
        output_type=OutputTypes.TEXT,
        tags=[Tags.CHATBOTS, Tags.CONTEXT, Tags.LLM, Tags.TEXT_COMPLETION, Tags.TEXT, Tags.STREAMING],
    )

    IS_STREAMING = True

    def init_llm_model(self) -> AsyncLLMEngine:
        """Initializes the vLLM async engine for real token-level streaming.

        Returns:
            AsyncLLMEngine: An initialized async vLLM engine.
        """
        engine_args = AsyncEngineArgs(
            model=self.attributes.init_args.llm_model_name,
            **self.attributes.init_args.model_dump(exclude={"llm_model_name"}),
        )
        return AsyncLLMEngine.from_engine_args(engine_args)

    def initialize(self) -> None:
        """Initializes the async engine, system prompt, and streaming sampling parameters.

        Extends the parent to set ``output_kind=RequestOutputKind.DELTA``
        so that each ``RequestOutput`` contains only new delta tokens.
        """
        super().initialize()
        self.sampling_params.output_kind = RequestOutputKind.DELTA

    async def _stream_responses(self, input_message: list[dict]) -> AsyncGenerator[str, None]:
        """Runs true token-level streaming via the vLLM async engine.

        Applies the model's chat template to render messages into a prompt,
        then iterates the ``AsyncLLMEngine.generate()`` async generator to
        yield each token delta as it is produced.

        Args:
            input_message (list[dict]): The formatted list of chat messages.

        Yields:
            str: Individual token chunks from the model as they are generated.
        """
        tokenizer = self.llm.get_tokenizer()
        prompt = tokenizer.apply_chat_template(
            input_message,
            tokenize=False,
            add_generation_prompt=True,
        )

        request_id = str(uuid.uuid4())

        async for output in self.llm.generate(prompt, self.sampling_params, request_id):
            for completion_output in output.outputs:
                delta = completion_output.text
                if delta:
                    self.logger.debug(f"Yielding chunk: '{delta}'")
                    yield delta

    async def get_response(self, input_message: list) -> AsyncGenerator[str, None]:
        """Asynchronously gets streaming response chunks, handling potential errors and retries.

        Args:
            input_message (Union[str, List]): The input messages list for the LLM.

        Yields:
            str: Text chunks from the streaming response.
        """
        self.logger.debug(f"Query is {input_message}")
        try:
            async for chunk in self._stream_responses(input_message):
                yield chunk
        except (IndexError, AttributeError):
            self.logger.warning("Error during streaming, retrying...")
            async for chunk in self._stream_responses(input_message):
                yield chunk

    async def generate_response(self, container: DataContainer) -> AsyncGenerator[DataContainer, None]:
        """Builds context and yields partial DataContainers for each response chunk.

        Args:
            container (DataContainer): The input container with text packets.

        Yields:
            DataContainer: Partial containers with one response chunk each, then the final container with EOS sentinel.
        """
        for packet in container.texts:
            full_context: list[dict] = []
            user_id, session_id, prompt = self.prepare_conversation_context(packet)
            if self.system_prompt:
                system_prompt_msg = self.generate_dict_msg(LLMChatKeys.system_value, self.system_prompt)
                full_context.append(system_prompt_msg)

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

            message = self.generate_dict_msg(LLMChatKeys.user_value, prompt)
            full_context.append(message)

            accumulated_text = ""
            async for response_chunk in self.get_response(full_context):
                if not response_chunk:
                    break

                accumulated_text += response_chunk
                processed_chunk = postprocess_text(response_chunk, self.attributes.pattern, self.attributes.keep_before)
                if not processed_chunk:
                    continue

                partial_container: DataContainer = deepcopy(container)
                partial_container.texts.append(TextPacket(source=session_id, content=processed_chunk, id=user_id))
                yield partial_container

        # add sentinel
        container.generic_data[SENTINEL_GENERIC_KEY] = SINAPSIS_END_OF_STREAM
        yield container

    async def async_execute(self, container: DataContainer) -> AsyncGenerator[DataContainer, None]:
        """Asynchronously executes the streaming generation process.

        Args:
            container (DataContainer): The input data container.

        Yields:
            DataContainer: Partial containers with response chunks, followed by
                           the final container with EOS sentinel.
        """
        try:
            async for partial_container in self.generate_response(container):
                yield partial_container
        except (IndexError, TypeError) as e:
            self.logger.error(f"got error when stopping {e}")

    def _shut_down(self) -> None:
        """Initiates a graceful shutdown of the vLLM async engine."""
        self.logger.info(f"Initiating shutdown for vLLM instance `{self.instance_name}`...")
        if hasattr(self, "llm") and self.llm is not None:
            self.llm.shutdown()
