import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import List, Optional, Dict, Any

class FinetunedModel:
    def __init__(self, model_path: str = "../finetuned_model"):
        """
        Initialize the finetuned model.
        
        Args:
            model_path: Path to the finetuned model directory
        """
        self.model_path = model_path
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        print(f"Loading finetuned model from {model_path}")
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForCausalLM.from_pretrained(model_path)
        self.model.to(self.device)
        
        # Set pad token if not set
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
    
    def generate(
        self,
        prompt: str,
        max_length: int = 200,
        num_return_sequences: int = 1,
        temperature: float = 0.7,
        top_p: float = 0.9,
        **kwargs
    ) -> List[str]:
        """
        Generate text using the finetuned model.
        
        Args:
            prompt: The prompt to generate text from
            max_length: Maximum length of the generated text
            num_return_sequences: Number of sequences to generate
            temperature: Sampling temperature (higher = more random)
            top_p: Nucleus sampling parameter
            **kwargs: Additional arguments to pass to the model
            
        Returns:
            List of generated text sequences
        """
        # Format the prompt
        formatted_prompt = f"Prompt: {prompt}\nCompletion:"
        
        # Tokenize the prompt
        inputs = self.tokenizer(formatted_prompt, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Generate text
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_length=max_length,
                num_return_sequences=num_return_sequences,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
                **kwargs
            )
        
        # Decode the generated text
        generated_texts = []
        for output in outputs:
            generated_text = self.tokenizer.decode(output, skip_special_tokens=True)
            # Extract just the completion part
            completion = generated_text.split("Completion:")[-1].strip()
            generated_texts.append(completion)
        
        return generated_texts
    
    def generate_with_context(
        self,
        prompt: str,
        content_type: str,
        similar_content: List[str],
        tone: Optional[str] = None,
        target_audience: Optional[str] = None,
        max_length: int = 200,
        num_return_sequences: int = 1,
        temperature: float = 0.7,
        top_p: float = 0.9,
        **kwargs
    ) -> List[str]:
        """
        Generate text with additional context.
        
        Args:
            prompt: The main prompt for content generation
            content_type: Type of content (email, social media, etc.)
            similar_content: List of similar content for context
            tone: Optional tone specification
            target_audience: Optional target audience specification
            max_length: Maximum length of the generated text
            num_return_sequences: Number of sequences to generate
            temperature: Sampling temperature (higher = more random)
            top_p: Nucleus sampling parameter
            **kwargs: Additional arguments to pass to the model
            
        Returns:
            List of generated text sequences
        """
        # Build the context
        context = self._build_context(prompt, content_type, similar_content, tone, target_audience)
        
        # Generate text
        return self.generate(
            prompt=context,
            max_length=max_length,
            num_return_sequences=num_return_sequences,
            temperature=temperature,
            top_p=top_p,
            **kwargs
        )
    
    def _build_context(
        self,
        prompt: str,
        content_type: str,
        similar_content: List[str],
        tone: Optional[str],
        target_audience: Optional[str]
    ) -> str:
        """
        Build a context string for the model.
        
        Args:
            prompt: The main prompt for content generation
            content_type: Type of content (email, social media, etc.)
            similar_content: List of similar content for context
            tone: Optional tone specification
            target_audience: Optional target audience specification
            
        Returns:
            Context string for the model
        """
        context = f"Content Type: {content_type}\n"
        if tone:
            context += f"Tone: {tone}\n"
        if target_audience:
            context += f"Target Audience: {target_audience}\n"
        
        context += "\nSimilar Content Examples:\n"
        for content in similar_content[:3]:  # Use top 3 similar content pieces
            context += f"- {content}\n"
        
        context += f"\nGenerate marketing copy for: {prompt}\n"
        return context

# Initialize the model
finetuned_model = FinetunedModel() 