from transformers import pipeline
from typing import List, Optional
import torch
from finetuned_model import finetuned_model

class MarketingCopywriter:
    def __init__(self):
        # Use the finetuned model instead of the default GPT-2
        self.model = finetuned_model
        
    def generate(
        self,
        prompt: str,
        content_type: str,
        similar_content: List[str],
        tone: Optional[str] = None,
    ) -> str:
        # Generate the marketing copy using the finetuned model
        generated_texts = self.model.generate_with_context(
            prompt=prompt,
            content_type=content_type,
            similar_content=similar_content,
            tone=tone,
            max_length=500,
            num_return_sequences=1,
            temperature=0.7,
            top_p=0.9
        )
        
        # Return the first generated text
        return generated_texts[0] if generated_texts else ""
    
    def _build_context(
        self,
        prompt: str,
        content_type: str,
        similar_content: List[str],
        tone: Optional[str],
        target_audience: Optional[str]
    ) -> str:
        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
    
    def _post_process(self, text: str) -> str:
        # Clean up the generated text
        text = text.strip()
        # Add any additional post-processing steps here
        return text

# Initialize the copywriter
copywriter = MarketingCopywriter()

def generate_marketing_copy(
    prompt: str,
    content_type: str,
    similar_content: List[str],
    tone: Optional[str] = None,
    target_audience: Optional[str] = None
) -> str:
    """
    Generate marketing copy based on the given parameters.
    
    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:
        Generated marketing copy
    """
    return copywriter.generate(
        prompt=prompt,
        content_type=content_type,
        similar_content=similar_content,
        tone=tone,
        target_audience=target_audience
    ) 


generate_marketing_copy(
    prompt="Help me write a blog post about the benefits of using our product",
    content_type="blog post",
    similar_content=[],
    tone="",
    target_audience=""
)