# Unsloth should be imported before transformers to ensure all optimizations are applied.
from unsloth import FastLanguageModel

from pathlib import Path
from threading import Thread
from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
from typing import Callable, Iterator, Optional
from xml_schema_validator import XmlLogitsProcessor
import os
import torch

from . import LlmEngine
from .. import util

class QwQLlmEngine(LlmEngine):
    
    def __init__(
        self,
        model_path: Path,
        temperature: float,
        xml_schema_text: Optional[str] = None,
    ):
        """
        Initialize the QwQ LLM Engine.

        Args:
            model_path: Local path to the model
            temperature: Sampling temperature
            xml_schema_text: Optional XML schema to validate against
        """
        self._temperature = temperature

        # Load tokenizer
        self._tokenizer = AutoTokenizer.from_pretrained(
            model_path,
        )
        
        # Load model
        self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
            model_path,
            load_in_4bit = True, # False for LoRA 16bit
            fast_inference = True, # Enable vLLM fast inference
            gpu_memory_utilization = 0.8, # Reduce if out of memory
            tokenizer = self._tokenizer,
        )
        
        # Create inference pipeline
        self._pipeline = pipeline(
            "text-generation",
            model=self._model,
            tokenizer=self._tokenizer,
            return_full_text=False,
        )

        if xml_schema_text:
            self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
        else:
            self._logits_processor = None

                        
    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
        """
        Run inference using the system prompt and main context.

        Args:
            system_prompt: The system prompt string
            main_context: The main context string after templating
            continuation_text: Part of the response that is already generated
            should_stop: Callback that returns True when inference should stop

        Returns:
            Iterator[str]: An iterator that yields the generated text.
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": main_context},
            {"role": "assistant", "content": continuation_text},
        ]

        text = self._tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False,
        )

        streamer = TextIteratorStreamer(
            self._tokenizer,
            skip_prompt=True,
        )

        generation_kwargs = {
            "text_inputs": text,
            "do_sample": True,
            "temperature": self._temperature,
            "streamer": streamer,
            "use_cache": True,
        }
        
        if self._logits_processor:
            generation_kwargs["logits_processor"] = [self._logits_processor.copy()]

        generation_thread = Thread(
            target=self._pipeline,
            kwargs=generation_kwargs
        )

        generation_thread.start()

        for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
            yield text
            if should_stop():
                break

        generation_thread.join()

    def token_count(self, system_prompt: str, main_context: str) -> int:
        """
        Count tokens for the given system prompt and main context.
        
        Args:
            system_prompt: The system prompt string
            main_context: The main context string
            
        Returns:
            int: Total number of tokens
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": main_context}
        ]
        prompt = self._tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        return len(self._tokenizer.encode(prompt))

    def token_limit(self) -> int:
        return self._pipeline.model.config.max_position_embeddings