from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
from typing import Iterator, Optional, Callable
from xml_schema_validator import XmlLogitsProcessor
import sys
import torch

from . import LlmEngine
from .. import util

class LocalLlmEngine(LlmEngine):
    def __init__(
        self,
        model_path: str,
        temperature: float,
        token_limit: int,
        xml_schema_text: Optional[str] = None,
        api_token: Optional[str] = None,
    ):
        """
        Initialize the LLM Engine with a model path.

        Args:
            model_path: Path to the model weights to be used.
            temperature: Temperature for sampling
            token_limit: Maximum number of tokens to generate
            xml_schema_text: Optional XML schema to validate against
            api_token: Huggingface API key
        """
        self._temperature = temperature
        self._token_limit = token_limit
        self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            return_dict=True,
            low_cpu_mem_usage=True,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            trust_remote_code=True,
            token=api_token,
        )
        if self._tokenizer.pad_token_id is None:
            self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
        if model.config.pad_token_id is None:
            model.config.pad_token_id = model.config.eos_token_id
        self._pipeline = pipeline(
            "text-generation",
            model=model,
            tokenizer=self._tokenizer,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            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, 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}
        ]
        prompt = self._tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        streamer = TextIteratorStreamer(
            self._tokenizer,
            skip_prompt=True
        )
        generation_kwargs = {
            "text_inputs": prompt,
            "do_sample": True,
            "temperature": self._temperature,
            "max_new_tokens": self.token_limit(),
            "streamer": streamer,
        }
        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:
        """
        Get the model's context window size.
        
        Returns:
            int: Maximum number of tokens the model can process
        """
        if self._token_limit is not None:
            return self._token_limit
        else:
            return self._pipeline.model.config.max_position_embeddings