import gradio as gr
import torch
from transformers import pipeline, set_seed
import re
import random
import os

# Set up the text generation pipeline with better error handling
def initialize_model():
    """Initialize the text generation model with fallback options"""
    try:
        # Check if CUDA is available and model can fit in memory
        device = 0 if torch.cuda.is_available() else -1
        print(f"Using device: {'GPU' if device == 0 else 'CPU'}")
        
        # Try GPT2-medium first, fallback to GPT2 if memory issues
        try:
            generator = pipeline(
                "text-generation",
                model="gpt2-medium",
                device=device,
                pad_token_id=50256,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
            )
            print("Loaded GPT2-medium model successfully")
        except Exception as e:
            print(f"GPT2-medium failed, falling back to GPT2: {e}")
            generator = pipeline(
                "text-generation",
                model="gpt2",
                device=device,
                pad_token_id=50256
            )
            print("Loaded GPT2 model successfully")
        
        return generator
    except Exception as e:
        print(f"Error loading model: {e}")
        # Return None to use fallback captions only
        return None

# Initialize model at startup
generator = initialize_model()

def is_valid_caption(caption):
    """Check if caption is valid and makes sense"""
    if not caption or len(caption.strip()) < 3:
        return False
    
    words = caption.split()
    
    # Check for unwanted patterns
    invalid_patterns = [
        r'http[s]?://',  # URLs
        r'www\.',        # Website references
        r'\.com',        # Domain extensions
        r'\.org',
        r'\.net',
        r'\d{3,}',       # Long numbers
        r'[^\w\s]',      # Special characters (except basic punctuation)
    ]
    
    for pattern in invalid_patterns:
        if re.search(pattern, caption, re.IGNORECASE):
            return False
    
    # Check for incomplete sentences (common AI issues)
    incomplete_endings = [
        'i get', 'i am', 'i will', 'i have', 'i do', 'i can',
        'you get', 'you are', 'you will', 'you have', 'you do',
        'we get', 'we are', 'we will', 'we have', 'we do',
        'they get', 'they are', 'they will', 'they have',
        'there is a', 'there are', 'this is a', 'that is',
        'it is a', 'here is', 'when you', 'if you',
        'the best', 'the most', 'the only', 'the first'
    ]
    
    caption_lower = caption.lower()
    for ending in incomplete_endings:
        if caption_lower.endswith(ending):
            return False
    
    # Check for nonsensical word combinations
    nonsensical_words = [
        'lorem', 'ipsum', 'dolor', 'amet', 'consectetur',
        'adipiscing', 'elit', 'sed', 'eiusmod', 'tempor'
    ]
    
    for word in nonsensical_words:
        if word in caption_lower:
            return False
    
    # Must have at least 2 different words
    if len(set(words)) < 2:
        return False
    
    # Check for proper word structure (no single letters except 'I' and 'a')
    for word in words:
        if len(word) == 1 and word.lower() not in ['i', 'a']:
            return False
    
    return True

def generate_short_caption(description, max_retries=5):
    """Generate a short, catchy caption (max 5 words) with better quality control"""
    if not description.strip() or generator is None:
        return ""
    
    for attempt in range(max_retries):
        try:
            # Set random seed for variety
            set_seed(random.randint(1, 10000))
            
            # More specific prompt templates for complete captions
            prompts = [
                f"Complete Instagram caption about {description[:40]}:",
                f"Short social media caption: {description[:40]}. Caption:",
                f"Trendy Instagram post about {description[:40]}:",
                f"Cool caption for {description[:40]}:",
                f"Fun social media post: {description[:40]}. Text:"
            ]
            
            prompt = random.choice(prompts)
            
            generated = generator(
                prompt,
                max_length=len(prompt.split()) + 12,  # Allow more space for complete thoughts
                num_return_sequences=3,  # Generate multiple options
                temperature=0.7,  # Slightly lower temperature for more coherent text
                do_sample=True,
                top_p=0.85,
                repetition_penalty=1.4,
                pad_token_id=50256,
                truncation=True
            )
            
            # Try each generated option
            for result in generated:
                full_text = result['generated_text']
                caption = full_text.replace(prompt, "").strip()
                
                if caption:
                    # Clean the caption
                    caption = re.sub(r'[^\w\s\']', '', caption)  # Keep apostrophes
                    caption = ' '.join(caption.split())  # Normalize whitespace
                    
                    # Try different word lengths (3-5 words)
                    for max_words in [5, 4, 3]:
                        words = caption.split()[:max_words]
                        if len(words) >= 2:
                            test_caption = " ".join(words)
                            if is_valid_caption(test_caption):
                                return test_caption
        
        except Exception as e:
            print(f"Generation error (attempt {attempt + 1}): {e}")
            continue
    
    return ""

def generate_instagram_captions(description):
    """Generate 5 short, catchy Instagram captions (max 5 words each) using cool captions"""
    if not description.strip():
        return "Please enter a description of your photo! 📸"
    
    # Massive database of cool modern captions organized by vibe
    cool_captions_db = {
        "general_vibes": [
            "Good vibes only today",
            "Living our best life",
            "Squad goals right here",
            "Making memories with favorites",
            "Caught in the moment",
            "Energy levels unmatched today",
            "Vibes on point always",
            "Good times roll forever",
            "Living for these moments",
            "Making magic happen daily",
            "Good company great memories",
            "Vibes immaculate as always",
            "Energy infectious spread everywhere",
            "Making waves together always",
            "Living life in color",
            "Good vibes attract everything",
            "Caught being absolutely legendary",
            "Creating our own magic",
            "Blessed with the best",
            "Dreams do come true",
            "Forever young at heart",
            "Happiness looks like this"
        ],
        "friendship_squad": [
            "Squad deeper than ocean",
            "Friends like these matter",
            "Squad tight bond stronger",
            "Good energy good people",
            "Squad up for success",
            "Squad goals achieved today",
            "Team work makes dreams",
            "Squad deeper than family",
            "Friends forever and always",
            "Squad goals right here",
            "Squad assembled ready always",
            "Squad till the end",
            "Squad stronger than ever",
            "Squad vibes unmatched today",
            "Squad goals achieved finally",
            "Squad energy off charts",
            "Squad love runs deep",
            "Squad goals never ending",
            "Squad till we die",
            "Squad brings out best"
        ],
        "confident_boss": [
            "Boss moves only today",
            "Confidence level maximum reached",
            "Main character energy activated",
            "That girl energy today",
            "Boss babe vibes only",
            "Confidence is my outfit",
            "Serving looks and attitude",
            "Boss energy never fades",
            "Confidence on ten always",
            "Boss mode permanently on",
            "Confidence through the roof",
            "Boss vibes all day",
            "Confidence is my superpower",
            "Boss lady energy strong",
            "Confidence level expert mode",
            "Boss moves make history",
            "Confidence speaks for itself",
            "Boss energy infectious spreading",
            "Confidence is my weapon",
            "Boss vibes never off"
        ],
        "aesthetic_mood": [
            "Aesthetic vibes only today",
            "Mood captured perfectly here",
            "Aesthetic game too strong",
            "Vibes aesthetic perfection achieved",
            "Mood board inspiration daily",
            "Aesthetic energy flowing freely",
            "Mood ring says happy",
            "Aesthetic goals achieved today",
            "Mood boost successfully delivered",
            "Aesthetic vibes never ending",
            "Mood elevated to maximum",
            "Aesthetic perfection in progress",
            "Mood lighting hits different",
            "Aesthetic energy through roof",
            "Mood captured in pixels",
            "Aesthetic vibes permanently on",
            "Mood board worthy moment",
            "Aesthetic goals finally achieved",
            "Mood enhanced by beauty",
            "Aesthetic vibes multiply daily"
        ],
        "fun_playful": [
            "Chaos and we prosper",
            "Weekend warriors in action",
            "Sunshine mixed with chaos",
            "Adventure awaits us always",
            "Mischief managed successfully today",
            "Chaos coordinator reporting duty",
            "Fun police cant stop",
            "Chaos with a purpose",
            "Mischief level expert achieved",
            "Chaos brings us together",
            "Fun mode permanently activated",
            "Chaos creates best memories",
            "Mischief makes life interesting",
            "Chaos with good intentions",
            "Fun times multiply exponentially",
            "Chaos theory in action",
            "Mischief planned and executed",
            "Chaos brings out creativity",
            "Fun vibes spreading everywhere",
            "Chaos with positive energy"
        ],
        "motivational_success": [
            "Grinding until dreams reality",
            "Success tastes so sweet",
            "Hustle hard dream bigger",
            "Grinding never stops here",
            "Success is only option",
            "Hustle mode permanently on",
            "Grinding toward bigger goals",
            "Success follows consistent effort",
            "Hustle beats talent consistently",
            "Grinding creates diamond pressure",
            "Success is earned daily",
            "Hustle culture lifestyle choice",
            "Grinding through every obstacle",
            "Success mindset never rests",
            "Hustle harder than yesterday",
            "Grinding until goals achieved",
            "Success requires relentless pursuit",
            "Hustle creates unstoppable momentum",
            "Grinding builds character strength",
            "Success is inevitable outcome"
        ],
        "chill_relaxed": [
            "Zen mode activated today",
            "Peace love and positivity",
            "Chill vibes only please",
            "Zen energy flowing freely",
            "Peace found in simplicity",
            "Chill mode permanently on",
            "Zen master in training",
            "Peace within creates harmony",
            "Chill vibes multiply naturally",
            "Zen state of mind",
            "Peace radiates from within",
            "Chill energy infectious spreading",
            "Zen garden of thoughts",
            "Peace is my superpower",
            "Chill vibes heal everything",
            "Zen moments create clarity",
            "Peace found in present",
            "Chill mode healing process",
            "Zen energy restores balance",
            "Peace multiplies when shared"
        ],
        "trendy_gen_z": [
            "No cap this slaps",
            "Periodt end of discussion",
            "This hits different though",
            "Absolutely serving looks today",
            "This is it chief",
            "We understood the assignment",
            "This absolutely sends me",
            "Living rent free here",
            "This is the vibe",
            "Absolutely no skips today",
            "This is so valid",
            "We really did that",
            "This is the moment",
            "Absolutely understood assignment perfectly",
            "This hits every time",
            "We said what said",
            "This is peak content",
            "Absolutely no notes needed",
            "This is the energy",
            "We chose violence today"
        ],
        "night_out": [
            "Night out squad assembled",
            "Party mode fully activated",
            "Night vibes hit different",
            "Party till sunrise today",
            "Night out energy unmatched",
            "Party like there tomorrow",
            "Night out goals achieved",
            "Party mode never ends",
            "Night vibes absolutely immaculate",
            "Party energy through roof",
            "Night out memories made",
            "Party until legs hurt",
            "Night out squad goals",
            "Party mode is lifestyle",
            "Night vibes multiply exponentially",
            "Party energy infectious spreading",
            "Night out therapy session",
            "Party mode permanently activated",
            "Night vibes heal everything",
            "Party like world watching"
        ],
        "self_love": [
            "Self love Sunday activated",
            "Main character energy strong",
            "Self care Sunday vibes",
            "Self love looks good",
            "Main character moment captured",
            "Self care never optional",
            "Self love is superpower",
            "Main character energy activated",
            "Self care Sunday ritual",
            "Self love journey continues",
            "Main character vibes only",
            "Self care is healthcare",
            "Self love wins always",
            "Main character development complete",
            "Self care Sunday success",
            "Self love radiates outward",
            "Main character energy infectious",
            "Self care heals everything",
            "Self love multiplies daily",
            "Main character story continues"
        ]
    }
    
    # Flatten all cool captions into one big list
    all_cool_captions = []
    for category, captions in cool_captions_db.items():
        all_cool_captions.extend(captions)
    
    # Generate AI captions if model is available
    ai_captions = []
    if generator is not None:
        attempts = 0
        while len(ai_captions) < 2 and attempts < 10:  # Try to get at least 2 good AI captions
            caption = generate_short_caption(description)
            if caption and caption not in ai_captions and is_valid_caption(caption):
                ai_captions.append(caption)
            attempts += 1
    
    # Combine AI and predefined cool captions
    combined_captions = ai_captions + all_cool_captions
    
    # Remove duplicates while preserving order
    seen = set()
    unique_captions = []
    for caption in combined_captions:
        if caption not in seen:
            seen.add(caption)
            unique_captions.append(caption)
    
    # Select 5 random captions
    selected_captions = random.sample(unique_captions, min(5, len(unique_captions)))
    
    # Ensure each caption is max 5 words and properly formatted
    final_captions = []
    for caption in selected_captions:
        words = caption.split()[:5]
        if len(words) >= 2:
            final_captions.append(" ".join(words))
    
    # Fill with defaults if needed
    while len(final_captions) < 5:
        fallback = random.choice(all_cool_captions)
        words = fallback.split()[:5]
        final_caption = " ".join(words)
        if final_caption not in final_captions:
            final_captions.append(final_caption)
    
    # Format as simple numbered list
    result = ""
    emojis = ["🎯", "⚡", "🌟", "🔥", "💫"]
    
    for i, caption in enumerate(final_captions[:5]):
        result += f"{emojis[i]} {caption}\n"
    
    return result.strip()

# Enhanced Instagram-inspired theme with custom CSS
instagram_css = """
/* Instagram color scheme and modern design */
:root {
    --ig-primary: #E4405F;
    --ig-secondary: #833AB4;
    --ig-tertiary: #F77737;
    --ig-quaternary: #FCAF45;
    --ig-blue: #405DE6;
    --ig-bg: #FAFAFA;
    --ig-text: #262626;
    --ig-border: #DBDBDB;
    --ig-light: #FFFFFF;
}

/* Main container styling */
.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    min-height: 100vh;
    font-family: 'Instagram Sans', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
}

/* Header styling - Modified to be regular black text */
.main-header {
    text-align: center;
    color: #000000 !important;
    font-size: 3em;
    font-weight: 400;
    margin-bottom: 0.5em;
    padding: 1rem 0;
    background: rgba(255, 255, 255, 0.9);
    border-radius: 15px;
    margin: 1rem;
    box-shadow: 0 4px 20px rgba(0,0,0,0.1);
}

.subtitle {
    text-align: center;
    color: white;
    font-size: 1.2em;
    margin-bottom: 2em;
    text-shadow: 0 1px 2px rgba(0,0,0,0.3);
    font-weight: 300;
}

/* Card styling */
.card {
    background: rgba(255, 255, 255, 0.95) !important;
    border-radius: 20px !important;
    box-shadow: 0 8px 32px rgba(0,0,0,0.1) !important;
    backdrop-filter: blur(10px) !important;
    border: 1px solid rgba(255,255,255,0.2) !important;
    padding: 2rem !important;
    margin: 1rem !important;
    transition: all 0.3s ease !important;
}

.card:hover {
    transform: translateY(-5px) !important;
    box-shadow: 0 12px 40px rgba(0,0,0,0.15) !important;
}

/* Input styling */
.gr-textbox {
    border-radius: 15px !important;
    border: 2px solid var(--ig-border) !important;
    transition: all 0.3s ease !important;
    font-size: 1.1em !important;
    padding: 1rem !important;
}

.gr-textbox:focus {
    border-color: var(--ig-primary) !important;
    box-shadow: 0 0 0 3px rgba(228, 64, 95, 0.1) !important;
}

/* Button styling */
.generate-btn {
    background: linear-gradient(45deg, var(--ig-primary), var(--ig-secondary)) !important;
    border: none !important;
    color: white !important;
    font-weight: bold !important;
    font-size: 1.2em !important;
    padding: 15px 30px !important;
    border-radius: 25px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(228, 64, 95, 0.3) !important;
    cursor: pointer !important;
}

.generate-btn:hover {
    transform: translateY(-3px) !important;
    box-shadow: 0 8px 25px rgba(228, 64, 95, 0.4) !important;
    background: linear-gradient(45deg, var(--ig-secondary), var(--ig-primary)) !important;
}

.generate-btn:active {
    transform: translateY(-1px) !important;
}

/* Output styling */
.output-box {
    background: rgba(255, 255, 255, 0.9) !important;
    border-radius: 15px !important;
    border: 2px solid var(--ig-border) !important;
    font-family: 'Instagram Sans', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
    font-size: 1.1em !important;
    line-height: 1.8 !important;
    padding: 1rem !important;
}

/* Examples styling */
.gr-examples {
    background: rgba(255, 255, 255, 0.8) !important;
    border-radius: 15px !important;
    padding: 1rem !important;
    margin-top: 1rem !important;
}

.gr-examples .gr-button {
    background: linear-gradient(45deg, var(--ig-tertiary), var(--ig-quaternary)) !important;
    border: none !important;
    border-radius: 10px !important;
    color: white !important;
    margin: 0.25rem !important;
    transition: all 0.3s ease !important;
    font-weight: 600 !important;
}

.gr-examples .gr-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 12px rgba(0,0,0,0.2) !important;
}

/* Tips section */
.tips-section {
    background: rgba(255, 255, 255, 0.9) !important;
    border-radius: 15px !important;
    padding: 1.5rem !important;
    margin-top: 2rem !important;
    text-align: center !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.1) !important;
}

.tips-section p {
    color: var(--ig-text) !important;
    font-size: 1em !important;
    margin: 0.5rem 0 !important;
}

/* Footer styling */
.footer {
    text-align: center;
    color: white;
    font-size: 0.9em;
    margin-top: 2rem;
    padding: 1rem;
    opacity: 0.8;
}

/* Mobile responsive */
@media (max-width: 768px) {
    .main-header {
        font-size: 2.2em !important;
    }
    
    .card {
        margin: 0.5rem !important;
        padding: 1.5rem !important;
    }
    
    .generate-btn {
        font-size: 1.1em !important;
        padding: 12px 25px !important;
    }
}

/* Loading animation */
@keyframes pulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.5; }
}

.loading {
    animation: pulse 2s infinite;
}
"""

# Create the Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="red",
        secondary_hue="purple",
        neutral_hue="gray"
    ),
    css=instagram_css,
    title="InstaCap - AI Caption Generator - By Yuval Alon Yoav Barnea and Dana Shpak"
) as demo:
    
    # Header
    gr.HTML('<h1 class="main-header">InstaCap – AI Caption Generator</h1>')
    gr.HTML('<p class="subtitle">By Yuval Alon Yoav Barnea and Dana Shpak</p>')
    gr.HTML('<p class="subtitle">Create short, catchy Instagram captions in seconds! Perfect for your social media posts ✨</p>')
    
    with gr.Row():
        with gr.Column(scale=1, elem_classes="card"):
            description_input = gr.Textbox(
                label="📝 Describe your photo",
                placeholder="E.g., group photo with friends at sunset, coffee date with bestie, hiking adventure with squad...",
                lines=3,
                max_lines=5,
                elem_classes="gr-textbox"
            )
            
            generate_btn = gr.Button(
                "✨ Generate Captions",
                variant="primary",
                elem_classes="generate-btn"
            )
        
        with gr.Column(scale=1, elem_classes="card"):
            caption_output = gr.Textbox(
                label="🎯 Your Instagram Captions",
                lines=8,
                max_lines=12,
                interactive=False,
                show_copy_button=True,
                elem_classes="output-box"
            )
    
    # Connect the button to the function
    generate_btn.click(
        fn=generate_instagram_captions,
        inputs=description_input,
        outputs=caption_output
    )
    
    # Also trigger on Enter key
    description_input.submit(
        fn=generate_instagram_captions,
        inputs=description_input,
        outputs=caption_output
    )
    
    gr.HTML(
        """
        <div class="tips-section">
            <p>💡 <strong>Pro Tips for Better Captions:</strong></p>
            <p>🎯 Each caption is maximum 5 words for maximum impact</p>
            <p>⚡ Get fresh trendy captions every time you generate</p>
            <p>📱 Copy and paste directly to Instagram, TikTok, or any social platform</p>
            <p>🔄 Click generate again for completely new cool options</p>
            <p>✨ Mix of AI-generated and trending cool captions</p>
        </div>
        """
    )
    
    gr.HTML(
        """
        <div class="footer">
            <p>Made with ❤️ for social media creators • Perfect for Instagram, TikTok, and more!</p>
        </div>
        """
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )