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

# Previous imports and model loading code remains the same...
# (Keep all the previous code until the create_interface function)

def create_interface():
    print("Loading models...")
    generator_tokenizer, generator, sentiment_analyzer, content_checker = load_models()
    print("Models loaded successfully!")
    
    # Sample data function
    def fill_sample_data():
        return [
            "EcoBottle",  # Product Name
            "Sustainable water bottle made from recycled ocean plastic",  # Product Description
            "Environmentally conscious young professionals",  # Target Audience
            "100% recycled materials, Insulated design, Leak-proof",  # Key Features
            "Helps clean oceans, Keeps drinks cold for 24 hours",  # Unique Benefits
            "Twitter",  # Platform
            "professional"  # Tone
        ]
    
    def process_input(
        product_name,
        product_description,
        target_audience,
        key_features,
        unique_benefits,
        platform,
        tone
    ):
        try:
            results = generate_content(
                product_name,
                product_description,
                target_audience,
                key_features,
                unique_benefits,
                platform,
                tone,
                generator_tokenizer,
                generator,
                sentiment_analyzer,
                content_checker
            )
            
            output = "🎯 Generated Marketing Content:\n\n"
            for i, content in enumerate(results, 1):
                output += f"Version {i}:\n"
                output += f"📝 Content: {content['text']}\n"
                output += f"😊 Sentiment: {content['sentiment']}\n"
                output += f"✅ Safety Score: {content['safety_score']}\n"
                output += "-" * 50 + "\n"
            
            return output
        except Exception as e:
            return f"An error occurred: {str(e)}"

    # Create input components
    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")
    
    # Output component
    output = gr.Textbox(label="Generated Content", lines=10)

    # Create the interface with custom layout
    iface = gr.Interface(
        fn=process_input,
        inputs=[
            product_name,
            product_description,
            target_audience,
            key_features,
            unique_benefits,
            platform,
            tone
        ],
        outputs=output,
        title="Ethimar - AI Marketing Content Generator",
        description="""Generate ethical marketing content with AI-powered insights.
        ⏳ Note: First generation might take 3-5 minutes due to model loading.
        Subsequent generations will be faster!""",
        theme="default",
        examples=[
            [
                "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"
            ]
        ]
    )

    # Add the sample data button with custom styling
    fill_button = gr.Button(
        "Fill the form with sample data",
        variant="primary",
        scale=1,
        size="sm"
    )
    
    # Connect the button to the fill_sample_data function
    fill_button.click(
        fn=fill_sample_data,
        outputs=[
            product_name,
            product_description,
            target_audience,
            key_features,
            unique_benefits,
            platform,
            tone
        ]
    )

    return iface

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