import logging
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from huggingface_hub import login
from .config import Config

logger = logging.getLogger(__name__)

class ModelManager:
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.tokenizer = None
        self.model = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        # Login to Hugging Face Hub
        if Config.HUGGING_FACE_TOKEN:
            logger.info("Logging in to Hugging Face Hub")
            login(token=Config.HUGGING_FACE_TOKEN)
        
        # Initialize tokenizer and model
        self._init_tokenizer()
        self._init_model()
        
    def _init_tokenizer(self):
        """Initialize the tokenizer."""
        try:
            logger.info(f"Loading tokenizer: {self.model_name}")
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                token=Config.HUGGING_FACE_TOKEN
            )
            # Ensure we have the necessary special tokens
            special_tokens = {
                'pad_token': '[PAD]',
                'eos_token': '</s>',
                'bos_token': '<s>'
            }
            self.tokenizer.add_special_tokens(special_tokens)
            logger.info("Tokenizer loaded successfully.")
        except Exception as e:
            logger.error(f"Error loading tokenizer: {str(e)}")
            raise
            
    def _init_model(self):
        """Initialize the model."""
        try:
            logger.info(f"Loading model: {self.model_name}")
            
            # Load model with CPU configuration
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                device_map={"": self.device},
                torch_dtype=torch.float32,  # Use float32 for CPU
                token=Config.HUGGING_FACE_TOKEN,
                low_cpu_mem_usage=True
            )
            # Resize embeddings to match tokenizer
            self.model.resize_token_embeddings(len(self.tokenizer))
            logger.info(f"Using device: {self.device}")
        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            raise

    def generate_text(self, prompt: str, max_new_tokens: int = 1024) -> str:
        """Generate text from prompt."""
        try:
            # Encode the prompt
            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}

            # Generate response
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=True,
                    temperature=Config.TEMPERATURE,
                    top_p=Config.TOP_P,
                    pad_token_id=self.tokenizer.pad_token_id,
                    eos_token_id=self.tokenizer.eos_token_id,
                )

            # Decode and return the generated text
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            # Extract only the generated part (remove the prompt)
            response = generated_text[len(prompt):].strip()
            
            return response
            
        except Exception as e:
            logger.error(f"Error generating text: {str(e)}")
            return """- Issues:
- Error generating code review
- Model inference failed

- Improvements:
- Please try again
- Check model configuration

- Best Practices:
- Ensure proper model setup
- Verify token permissions

- Security:
- No immediate concerns"""