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

from . import LlmEngine
from .. import util

class DeepSeekLlmEngine(LlmEngine):
    """
    LLM Engine implementation for DeepSeek models.
    Supports fine-tuned DeepSeek-R1 and its distilled versions.
    """
    
    def __init__(
        self,
        model_path: str,
        temperature: float = 0.6,
        token_limit: Optional[int] = None,
        api_key: Optional[str] = None,
    ):
        """
        Initialize the DeepSeek LLM Engine.

        Args:
            model_path: Local path to the fine-tuned model
            temperature: Sampling temperature (0.6 default as recommended)
            token_limit: Maximum tokens to generate or context length override
            api_key: HuggingFace API token if needed
        """
        self._model_path = Path(model_path)
        self._temperature = temperature
        self._token_limit = token_limit
        
        # Load tokenizer with trust_remote_code for DeepSeek models
        self._tokenizer = AutoTokenizer.from_pretrained(
            self._model_path,
            token=api_key,
            trust_remote_code=True,
        )
        
        # Set padding token to avoid warnings
        if self._tokenizer.pad_token is None:
            self._tokenizer.pad_token = self._tokenizer.eos_token
        
        # Configure 4-bit quantization with CPU offloading
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            llm_int8_enable_fp32_cpu_offload=True
        )
        
        # Configure device map for efficient memory usage
        # "auto" with the proper quantization config will handle the memory constraints
        self._device_map = "auto"
        
        # Load model with quantization config
        self._model = AutoModelForCausalLM.from_pretrained(
            self._model_path,
            return_dict=True,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
            device_map=self._device_map,
            quantization_config=quantization_config,
            torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
            token=api_key,
        )
        
        # Ensure model is in evaluation mode
        self._model.eval()
                        
    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.
        """
        # Tokenize input
        inputs = self._tokenizer(system_prompt + "\n\n" + main_context, return_tensors="pt").to(self._device_map)
        
        # Create streamer for token-by-token generation
        streamer = TextIteratorStreamer(
            self._tokenizer,
            skip_prompt=True,
            timeout=15.0
        )
        
        # Generate in a separate thread to enable streaming
        generation_kwargs = {
            "input_ids": inputs.input_ids,
            "attention_mask": inputs.attention_mask,
            "max_new_tokens": self.token_limit() if self._token_limit else 2048,
            "temperature": self._temperature,
            "do_sample": True,
            "streamer": streamer,
            "repetition_penalty": 1.1,
            "pad_token_id": self._tokenizer.pad_token_id,
        }
        
        generation_thread = Thread(target=self._model.generate, kwargs=generation_kwargs)
        generation_thread.start()
        
        # Yield tokens as they become available
        try:
            for text in streamer:
                yield text
                if should_stop():
                    break
        finally:
            # Ensure thread is properly joined even if iteration is interrupted
            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
        """
        combined_prompt = f"{system_prompt}\n\n{main_context}"
        return len(self._tokenizer.encode(combined_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
        
        # Try to detect model size from config
        try:
            config_file = self._model_path / "config.json"
            if config_file.exists():
                import json
                with open(config_file, 'r') as f:
                    config = json.load(f)
                if 'max_position_embeddings' in config:
                    return config['max_position_embeddings']
                if 'model_max_length' in config:
                    return config['model_max_length']
        except Exception:
            pass
            
        # Default to 8k if we can't determine
        return 8192