import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

class TextGenerator:
    def __init__(self):
        print("Initializing Text Generator...")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")
        
        # Load model and tokenizer
        self.model_name = "facebook/opt-350m"
        print(f"Loading model {self.model_name}...")
        
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
        ).to(self.device)
        
        print(f"Model loaded and moved to {self.device}")

    def generate_text(self, prompt, max_length=200, temperature=0.7, top_p=0.9):
        """
        Generate text based on the given prompt
        
        Args:
            prompt (str): The text generation prompt
            max_length (int): Maximum length of the generated text
            temperature (float): Controls randomness in generation
            top_p (float): Controls diversity of generation
            
        Returns:
            str: Generated text
        """
        try:
            print(f"Generating text on {self.device}...")
            
            # Format prompt for better generation
            formatted_prompt = f"Instruction: {prompt}\n\nResponse:"
            
            inputs = self.tokenizer(
                formatted_prompt,
                return_tensors="pt",
                truncation=True,
                max_length=512
            ).to(self.device)
            
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_length=max_length + len(inputs["input_ids"][0]),
                    temperature=temperature,
                    top_p=top_p,
                    num_return_sequences=1,
                    pad_token_id=self.tokenizer.eos_token_id,
                    do_sample=True,
                    repetition_penalty=1.2,
                    no_repeat_ngram_size=3,
                    num_beams=5,
                    early_stopping=True
                )
            
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            # Remove the prompt from the generated text
            generated_text = generated_text[len(formatted_prompt):]
            
            # Format the text
            formatted_text = self._format_text(generated_text)
            return formatted_text

        except Exception as e:
            return f"Error generating text: {str(e)}"

    def _format_text(self, text):
        """
        Format the generated text for better readability
        
        Args:
            text (str): The text to format
            
        Returns:
            str: Formatted text
        """
        # Split into paragraphs
        paragraphs = text.split('\n\n')
        
        # Format each paragraph
        formatted_paragraphs = []
        for para in paragraphs:
            if para.strip():
                # Capitalize first letter
                para = para.strip()
                if para:
                    para = para[0].upper() + para[1:]
                
                # Add proper spacing
                para = ' '.join(para.split())
                
                formatted_paragraphs.append(para)
        
        # Join paragraphs with proper spacing
        formatted_text = '\n\n'.join(formatted_paragraphs)
        
        # Ensure proper punctuation
        if formatted_text and formatted_text[-1] not in '.!?':
            formatted_text += '.'
        
        return formatted_text