import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import nltk
from datetime import datetime, timedelta
import requests
from bs4 import BeautifulSoup

# Download required NLTK data
try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt')

# Global variables to cache models
CACHED_MODELS = {}

# Initialize models and tokenizers
def load_models():
    global CACHED_MODELS
    
    if CACHED_MODELS:
        return (
            CACHED_MODELS['generator_tokenizer'],
            CACHED_MODELS['generator'],
            CACHED_MODELS['sentiment_analyzer'],
            CACHED_MODELS['content_checker']
        )
    
    try:
        # Use GPT-2 instead of DistilGPT-2
        generator_model = "gpt2"
        generator_tokenizer = AutoTokenizer.from_pretrained(generator_model)
        generator = AutoModelForCausalLM.from_pretrained(generator_model)
        
        # Sentiment analysis
        sentiment_analyzer = pipeline(
            "sentiment-analysis",
            model="finiteautomata/bertweet-base-sentiment-analysis"
        )
        
        # Content safety checker
        content_checker = pipeline(
            "text-classification",
            model="facebook/roberta-hate-speech-dynabench-r4-target"
        )
        
        # Cache the models
        CACHED_MODELS['generator_tokenizer'] = generator_tokenizer
        CACHED_MODELS['generator'] = generator
        CACHED_MODELS['sentiment_analyzer'] = sentiment_analyzer
        CACHED_MODELS['content_checker'] = content_checker
        
        return generator_tokenizer, generator, sentiment_analyzer, content_checker
    except Exception as e:
        print(f"Error loading models: {str(e)}")
        raise

def generate_content(
    product_name,
    product_description,
    target_audience,
    key_features,
    unique_benefits,
    platform,
    tone,
    generator_tokenizer,
    generator,
    sentiment_analyzer,
    content_checker
):
    char_limit = 280 if platform == "Twitter" else 500
    
    # Split features and benefits into lists
    features = [f.strip() for f in key_features.split(',')]
    benefits = [b.strip() for b in unique_benefits.split(',')]
    
    # Expanded intro phrases for each tone
    intro_phrases = {
        'professional': [
            f"Introducing {product_name}:",
            f"Discover {product_name}:",
            f"Meet {product_name}:",
            f"Presenting {product_name}:",
            f"Announcing {product_name}:",
            f"Experience {product_name}:",
            f"Elevate your life with {product_name}:",
        ],
        'casual': [
            f"Check out {product_name}!",
            f"Say hello to {product_name}!",
            f"Get ready for {product_name}!",
            f"You'll love {product_name}!",
            f"Looking for something special? Try {product_name}!",
            f"Meet your new favorite: {product_name}!",
            f"Game-changer alert: {product_name} is here!",
        ],
        'friendly': [
            f"We're excited to share {product_name} with you!",
            f"You'll love what {product_name} can do!",
            f"Let {product_name} transform your day!",
            f"Ready to discover {product_name}?",
            f"Here's why you'll love {product_name}:",
            f"Make every day better with {product_name}!",
            f"Transform your experience with {product_name}!",
        ]
    }

    # Description connectors
    description_connectors = [
        f" - {product_description}",
        f": {product_description}",
        f"! {product_description}",
        f", {product_description}",
        f". {product_description}",
    ]

    # Feature introduction phrases
    feature_intros = [
        "Featuring",
        "With",
        "Including",
        "Equipped with",
        "Powered by",
        "Designed with",
        "Built with",
        "Offering",
    ]

    # Benefit connectors
    benefit_connectors = [
        "Experience",
        "Enjoy",
        "Benefit from",
        "Take advantage of",
        "Discover",
        "Appreciate",
        "Make the most of",
    ]

    # Target audience phrases
    audience_phrases = [
        f"Perfect for {target_audience}",
        f"Ideal for {target_audience}",
        f"Designed for {target_audience}",
        f"Made specially for {target_audience}",
        f"Tailored for {target_audience}",
        f"Created with {target_audience} in mind",
    ]

    # Call-to-action phrases
    cta_phrases = {
        'Twitter': [
            "Learn more today!",
            "Discover more →",
            "Get yours now!",
            "Visit our website!",
            "Join us today!",
            "Transform your life today!",
            "Start your journey!",
        ],
        'Instagram': [
            f"\n\n#{product_name.replace(' ', '')}",
            f"\n\nLearn more - Link in bio! #{product_name.replace(' ', '')}",
            f"\n\nDiscover more ↗️ #{product_name.replace(' ', '')}",
            f"\n\nTap link to learn more! #{product_name.replace(' ', '')}",
        ]
    }

    import random

    def create_post():
        # Select tone
        selected_tone = tone.lower() if tone.lower() in intro_phrases else 'professional'
        
        # Random structure selection (1-4)
        structure = random.randint(1, 4)
        
        if structure == 1:
            # Standard structure
            post = random.choice(intro_phrases[selected_tone])
            post += random.choice(description_connectors)
            feature = random.choice(features)
            benefit = random.choice(benefits)
            
            if len(post) + len(feature) + len(benefit) + 20 < char_limit:
                post += f" {random.choice(feature_intros)} {feature}."
                post += f" {random.choice(benefit_connectors)} {benefit}."
                
        elif structure == 2:
            # Benefit-first structure
            benefit = random.choice(benefits)
            post = f"Ready to {benefit.lower()}? "
            post += random.choice(intro_phrases[selected_tone])
            post += random.choice(description_connectors)
            
            if len(post) + 30 < char_limit:
                feature = random.choice(features)
                post += f" {random.choice(feature_intros)} {feature}."
                
        elif structure == 3:
            # Question-based structure
            post = f"Looking for {product_description.lower()}? "
            post += random.choice(intro_phrases[selected_tone]).replace(':', '!')
            feature = random.choice(features)
            benefit = random.choice(benefits)
            
            if len(post) + len(feature) + len(benefit) + 20 < char_limit:
                post += f" {random.choice(feature_intros)} {feature}."
                post += f" {benefit}."
                
        else:
            # Feature-focused structure
            feature = random.choice(features)
            post = f"From {feature} to {random.choice(benefits).lower()}, "
            post += f"{product_name} has it all! "
            post += product_description

        # Add target audience if space allows
        if len(post) + 50 < char_limit:
            post += f" {random.choice(audience_phrases)}."

        # Add call to action
        if platform == "Twitter":
            if len(post) + 30 < char_limit:
                post += f" {random.choice(cta_phrases['Twitter'])}"
        else:
            if len(post) + 50 < char_limit:
                post += random.choice(cta_phrases['Instagram'])

        return post.strip()

    try:
        # Generate multiple versions
        posts = [create_post() for _ in range(2)]
        filtered_content = []
        
        for post in posts:
            # Verify length
            if len(post) > char_limit:
                post = post[:char_limit-3] + "..."
            
            # Check sentiment and safety
            try:
                sentiment = sentiment_analyzer(post)[0]
                safety_check = content_checker(post)[0]
                
                filtered_content.append({
                    'text': post,
                    'sentiment': sentiment['label'],
                    'safety_score': f"{float(safety_check.get('score', 0)):.2f}"
                })
            except Exception as e:
                print(f"Error in content analysis: {str(e)}")
                continue
        
        return filtered_content if filtered_content else [{
            'text': create_post(),
            'sentiment': 'positive',
            'safety_score': '1.00'
        }]
    
    except Exception as e:
        print(f"Error in content generation: {str(e)}")
        return [{
            'text': f"Introducing {product_name}: {product_description[:100]}... Learn more!",
            'sentiment': 'neutral',
            'safety_score': '1.00'
        }]

def create_interface():
    print("Loading models...")
    generator_tokenizer, generator, sentiment_analyzer, content_checker = load_models()
    print("Models loaded successfully!")
    
    def fill_sample_data():
        return [
            "EcoBottle",
            "Sustainable water bottle made from recycled ocean plastic",
            "Environmentally conscious young professionals",
            "100% recycled materials, Insulated design, Leak-proof",
            "Helps clean oceans, Keeps drinks cold for 24 hours",
            "Twitter",
            "professional"
        ]
    
    def clear_form():
        return [""] * 7  # Returns empty strings for all 7 input fields
    
    # Rest of your process_input_with_loading function remains the same...

    # Create the interface with blocks for custom layout
    with gr.Blocks(theme=gr.themes.Default()) as demo:
        gr.Markdown("# Ethimar - AI Marketing Content Generator")
        gr.Markdown("Generate ethical marketing content with AI-powered insights.\n⏳ Note: First generation might take 3-5 minutes due to model loading. Subsequent generations will be faster!")
        
        # Sample data button with custom styling
        with gr.Row():
            fill_button = gr.Button(
                "Fill the form with sample data",
                variant="primary",
                size="sm",
                scale=1
            )

        # Main content area with two columns
        with gr.Row():
            # Left column - Input form
            with gr.Column(scale=1):
                product_name = gr.Textbox(label="Product Name", placeholder="Enter product name")
                product_description = gr.Textbox(label="Product Description", lines=3, placeholder="Brief description of your product")
                target_audience = gr.Textbox(label="Target Audience", placeholder="Who is this product for?")
                key_features = gr.Textbox(label="Key Features", lines=2, placeholder="Main features of your product")
                unique_benefits = gr.Textbox(label="Unique Benefits", lines=2, placeholder="What makes your product special?")
                platform = gr.Radio(
                    choices=["Twitter", "Instagram"],
                    label="Platform",
                    value="Twitter"
                )
                tone = gr.Textbox(label="Tone", placeholder="e.g., professional, casual, friendly")
                
                # Buttons row at the bottom of the form
                with gr.Row():
                    submit_button = gr.Button("Generate Content", variant="primary", scale=2)
                    clear_button = gr.Button("Clear Form", variant="secondary", scale=1)

            # Right column - Output
            with gr.Column(scale=1):
                output = gr.Textbox(
                    label="Generated Content", 
                    lines=12,
                    value="✨ Enter your product details and click 'Generate Content' to start!"
                )

        # Connect all buttons to functions
        input_components = [
            product_name,
            product_description,
            target_audience,
            key_features,
            unique_benefits,
            platform,
            tone
        ]
        
        fill_button.click(
            fn=fill_sample_data,
            outputs=input_components
        )
        
        submit_button.click(
            fn=process_input_with_loading,
            inputs=input_components,
            outputs=output
        )
        
        clear_button.click(
            fn=clear_form,
            outputs=input_components
        )

    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_interface()
    demo.launch()