import json
from typing import Any

from pydantic import Field
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_chatbots_base.templates.llm_text_completion_base import (
    LLMTextCompletionAttributes,
    LLMTextCompletionBase,
)
from sinapsis_core.data_containers.data_packet import DataContainer, TextPacket
from sinapsis_core.template_base.base_models import OutputTypes, UIPropertiesMetadata
from vllm import LLM, SamplingParams
from vllm.distributed.parallel_state import destroy_model_parallel

from sinapsis_vllm.helpers.schemas import vLLMCompletionArgs, vLLMInitArgs


class vLLMTextCompletionAttributes(LLMTextCompletionAttributes):
    """Attributes for vLLM text completion templates.

    Inherits shared chatbot attributes from ``LLMTextCompletionAttributes``
    and adds vLLM-specific engine and completion arguments.

    Attributes:
        init_args (vLLMInitArgs): vLLM engine configuration arguments.
        completion_args (vLLMCompletionArgs): vLLM sampling and generation arguments.
            Defaults to ``vLLMCompletionArgs()``.
        structured_output_key (str): Key for storing parsed JSON output in
            the packet's ``generic_data``. Defaults to ``"structured_output"``.
    """

    init_args: vLLMInitArgs
    completion_args: vLLMCompletionArgs = Field(default_factory=vLLMCompletionArgs)
    structured_output_key: str = "structured_output"


class vLLMTextCompletion(LLMTextCompletionBase):
    """Template for configuring and running a vLLM-based text completion model.

    Usage example:

    agent:
      name: my_test_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: vLLMTextCompletion
      class_name: vLLMTextCompletion
      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
        chat_history_key: null
        rag_context_key: null
        system_prompt: You are an expert in AI.
        pattern: null
        keep_before: true
    """

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

    AttributesBaseModel = vLLMTextCompletionAttributes

    def init_llm_model(self) -> LLM:
        """Initializes the vLLM engine using the configuration attributes.

        This method creates a new vLLM ``LLM`` instance configured with the
        model name and all engine parameters from ``init_args``.

        Returns:
            LLM: An initialized instance of the vLLM engine.
        """
        return LLM(
            model=self.attributes.init_args.llm_model_name,
            **self.attributes.init_args.model_dump(exclude={"llm_model_name"}),
        )

    def initialize(self) -> None:
        """Initializes the LLM model, system prompt, and sampling parameters.

        Extends the base class initialization by also building ``SamplingParams``
        from the completion arguments. When a structured output response format
        is configured, the corresponding ``StructuredOutputsParams`` are injected
        into the sampling parameters.
        """
        super().initialize()
        completion_kwargs = self.attributes.completion_args.model_dump(exclude_none=True, exclude={"response_format"})
        structured_outputs = self.attributes.completion_args.response_format.to_vllm_format()
        if structured_outputs:
            completion_kwargs["structured_outputs"] = structured_outputs
        self.sampling_params = SamplingParams(**completion_kwargs)

    def get_response(self, input_message: str | list) -> str | None:
        """Generates a response from the vLLM model based on the provided message list.

        This method sends the input messages to the vLLM engine via ``llm.chat()``
        and extracts the generated text from the first output.

        Args:
            input_message (str | list): The input message list for chat completion.

        Returns:
            str | None: The model's response as a string, or None if no response
                is generated.
        """
        self.logger.debug(f"Query is {input_message}")
        outputs = self.llm.chat(
            messages=input_message,
            sampling_params=self.sampling_params,
            use_tqdm=False,
        )
        if outputs:
            response = outputs[0].outputs[0].text
            if response:
                return postprocess_text(response, self.attributes.pattern, self.attributes.keep_before)
        return None

    def _parse_json_response(self, response: str) -> dict[str, Any] | None:
        """Attempts to parse a JSON response string into a dictionary.

        Args:
            response (str): The raw response string to parse.

        Returns:
            dict[str, Any] | None: The parsed dictionary if valid JSON, None otherwise.
        """
        try:
            parsed = json.loads(response)
            if isinstance(parsed, dict):
                return parsed
        except (json.JSONDecodeError, TypeError):
            self.logger.warning(f"Failed to parse structured output as JSON: {response[:200]}")
        return None

    def generate_response(self, container: DataContainer) -> DataContainer:
        """Processes text packets and generates responses, parsing JSON for structured outputs.

        Extends the base class behavior by parsing JSON responses into the packet's
        generic_data when response_format type is ``'json_object'``. The raw string is
        always preserved as the content for compatibility with TextPacket.

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

        Returns:
            DataContainer: Updated DataContainer with the LLM responses.
        """
        self.logger.debug("Chatbot in progress")
        responses = []
        for packet in container.texts:
            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))
            response = self.infer(full_context)
            self.logger.debug("End of interaction.")

            if response:
                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))

        container.texts.extend(responses)
        return container

    def _shut_down(self) -> None:
        """Initiates a graceful shutdown of the vLLM engine.

        Clears the model, destroys the parallel state, and frees GPU memory.
        """
        self.logger.info(f"Initiating shutdown for vLLM instance `{self.instance_name}`...")
        self.clear_model()
        destroy_model_parallel()
        self.clear_memory()
