from pathlib import Path
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
from typing import Callable, Iterator
import json
import torch

from . import LlmEngine
from .. import util

class QwQLlmEngine(LlmEngine):
    
    def __init__(
        self,
        model_path: Path,
        temperature: float,
        token_limit: int = None,
    ):
        """
        Initialize the QwQ LLM Engine.

        Args:
            model_path: Local path to the model
            temperature: Sampling temperature
            token_limit: Maximum tokens to generate
        """
        self._temperature = temperature
        self._token_limit = token_limit

        with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
            tokenizer_config = json.load(f)

        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True
        )

        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            return_dict=True,
            device_map="auto",
            use_cache=True,
            quantization_config=quantization_config,
        )

        self._tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            tokenizer_config=tokenizer_config,
        )

        self._pipline = pipeline(
            "text-generation",
            model=model,
            tokenizer=self._tokenizer,
            return_full_text=False,
        )

                        
    def infer(self, system_prompt: str, main_context: 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
            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}
        ]

        text = self._tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
            
        streamer = TextIteratorStreamer(
            self._tokenizer,
            skip_prompt=True,
        )

        generation_thread = Thread(
            target=self._pipline,
            kwargs=dict(
                text_inputs=text,
                do_sample=True,
                temperature=self._temperature,
                max_new_tokens=self._token_limit,
                streamer=streamer,
            )
        )
            
        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._token_limit