"""
Local LLM Provider
==================

Local LLM inference without API calls using transformers or llama.cpp.
Supports CPU and GPU inference.
"""

from __future__ import annotations

import logging
from typing import Optional
import torch

logger = logging.getLogger(__name__)

# Check available backends
try:
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, pipeline
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False
    logger.warning("transformers not available. Install: pip install transformers torch")


class LocalLLM:
    """
    Local LLM for text polishing and correction.
    
    Supports multiple model types:
    - google/mt5-small: Multilingual T5 (good for Hindi, 300MB)
    - google/mt5-base: Larger mT5 (600MB, better quality)
    - IndicBART: Specialized for Indian languages
    """
    
    # Recommended models by size and quality
    RECOMMENDED_MODELS = {
        "small": "google/mt5-small",  # ~300MB, fast on CPU
        "base": "google/mt5-base",    # ~600MB, better quality
        "large": "google/mt5-large",  # ~1.2GB, best quality
    }
    
    def __init__(
        self,
        model_name: str = "google/mt5-small",
        device: str = "cpu",
        compute_type: str = "float32",
    ):
        """
        Initialize local LLM.
        
        Args:
            model_name: Model name or size ("small", "base", "large")
            device: "cpu" or "cuda"
            compute_type: "float32", "float16", or "int8"
        """
        if not TRANSFORMERS_AVAILABLE:
            raise RuntimeError("transformers not installed. Run: pip install transformers torch")
        
        # Resolve model name
        if model_name in self.RECOMMENDED_MODELS:
            model_name = self.RECOMMENDED_MODELS[model_name]
        
        self.model_name = model_name
        self.device = device
        self.compute_type = compute_type
        self.model = None
        self.tokenizer = None
        self._is_loaded = False
        
        logger.info(f"Initialized LocalLLM: {model_name} on {device}")
    
    def load(self):
        """Load model and tokenizer."""
        if self._is_loaded:
            return
        
        logger.info(f"Loading local LLM: {self.model_name}...")
        
        try:
            # Load tokenizer
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            
            # Determine torch dtype
            dtype = torch.float32
            if self.compute_type == "float16" and self.device == "cuda":
                dtype = torch.float16
            elif self.compute_type == "int8":
                dtype = torch.int8
            
            # Load model (T5 models are seq2seq)
            if "t5" in self.model_name.lower():
                self.model = AutoModelForSeq2SeqLM.from_pretrained(
                    self.model_name,
                    torch_dtype=dtype,
                    low_cpu_mem_usage=True,
                )
            else:
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_name,
                    torch_dtype=dtype,
                    low_cpu_mem_usage=True,
                )
            
            # Move to device
            if self.device == "cuda" and torch.cuda.is_available():
                self.model = self.model.to("cuda")
                logger.info("✅ LLM loaded on GPU")
            else:
                self.model = self.model.to("cpu")
                logger.info("✅ LLM loaded on CPU")
            
            self.model.eval()  # Set to evaluation mode
            self._is_loaded = True
            
        except Exception as e:
            logger.error(f"Failed to load local LLM: {e}")
            raise
    
    def generate(
        self,
        prompt: str,
        max_length: int = 512,
        temperature: float = 0.1,
        num_beams: int = 4,
    ) -> str:
        """
        Generate text using local LLM.
        
        Args:
            prompt: Input prompt/text to polish
            max_length: Maximum output length
            temperature: Sampling temperature (lower = more deterministic)
            num_beams: Number of beams for beam search
            
        Returns:
            Generated text
        """
        if not self._is_loaded:
            self.load()
        
        try:
            # Tokenize input
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                max_length=512,
                truncation=True,
                padding=True,
            )
            
            # Move to device
            if self.device == "cuda":
                inputs = {k: v.to("cuda") for k, v in inputs.items()}
            
            # Generate
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_length=max_length,
                    temperature=temperature,
                    num_beams=num_beams,
                    early_stopping=True,
                    do_sample=temperature > 0.0,
                )
            
            # Decode output
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            return generated_text.strip()
            
        except Exception as e:
            logger.error(f"Generation failed: {e}")
            return prompt  # Return original on error
    
    def polish_hindi_text(self, text: str) -> str:
        """
        Polish Hindi/Hinglish text with proper prompt.
        
        Args:
            text: Input Hindi/Hinglish text
            
        Returns:
            Polished text
        """
        # For T5 models, use task prefix
        if "t5" in self.model_name.lower():
            prompt = f"grammar: {text}"
        else:
            prompt = f"Correct spelling and grammar in this Hindi/Hinglish text: {text}\n\nCorrected:"
        
        return self.generate(prompt, max_length=len(text.split()) * 3)
    
    def is_available(self) -> bool:
        """Check if model is loaded and available."""
        return self._is_loaded
    
    def unload(self):
        """Unload model to free memory."""
        if self._is_loaded:
            del self.model
            del self.tokenizer
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            self._is_loaded = False
            logger.info("LLM unloaded from memory")


# Singleton instance
_local_llm_instance: Optional[LocalLLM] = None


def get_local_llm(
    model_name: str = "small",
    device: str = "cpu",
    compute_type: str = "float32",
) -> LocalLLM:
    """
    Get or create local LLM singleton instance.
    
    Args:
        model_name: Model name or size ("small", "base", "large")
        device: "cpu" or "cuda"
        compute_type: "float32", "float16", or "int8"
        
    Returns:
        LocalLLM instance
    """
    global _local_llm_instance
    
    if _local_llm_instance is None:
        _local_llm_instance = LocalLLM(
            model_name=model_name,
            device=device,
            compute_type=compute_type,
        )
        _local_llm_instance.load()
    
    return _local_llm_instance
