app.py

import streamlit as st
import google.generativeai as genai
from PIL import Image, ImageDraw, ImageFont
import io
import os
import base64
from typing import Optional

# Configure page
st.set_page_config(
    page_title="AI Meme Poster Creator",
    page_icon="🎭",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state
if 'generated_image' not in st.session_state:
    st.session_state.generated_image = None

class MemeGenerator:
    def __init__(self):
        try:
            from config import GEMINI_API_KEY
            self.api_key = GEMINI_API_KEY
        except ImportError:
            self.api_key = "YOUR_GEMINI_API_KEY_HERE"
        self.setup_gemini()

    def setup_gemini(self):
        """Setup Gemini AI configuration"""
        if self.api_key and self.api_key != "YOUR_GEMINI_API_KEY_HERE":
            try:
                genai.configure(api_key=self.api_key)
                self.model = genai.GenerativeModel('gemini-pro')
                st.sidebar.success("✅ AI Connected!")
            except Exception as e:
                st.sidebar.error(f"❌ API Error: {e}")
                self.model = None
        else:
            self.model = None
            st.sidebar.warning("⚠️ Add Gemini API key in config.py!")

    def generate_caption(self, theme: str, style: str = "funny") -> str:
        """ULTRA-RELIABLE CAPTION GENERATOR"""
        
        # EXTENSIVE MANUAL CAPTION DATABASE
        CAPTION_DB = {
            "college life": [
                "When the syllabus says 'easy A' but the first assignment says otherwise",
                "My sleep schedule: more messed up than my sleep schedule",
                "Coffee isn't a beverage, it's an emotional support liquid",
                "Me: I'll be productive today! Also me: watches 7 hours of cat videos",
                "Brain: store important info. Also brain: remembers every Vine from 2016",
                "Student budget: when ramen noodles feel like fine dining",
                "1% battery, 99% anxiety during online exams",
                "Sprinting across campus in 2 minutes flat: Olympic sport we didn't ask for",
                "Eating cold pizza for breakfast: the college gourmet experience",
                "Backpack weight: 50% books, 50% existential dread"
            ],
            "exams": [
                "When you study everything except what's on the exam",
                "That moment when you realize you brought the wrong calculator",
                "Multiple choice: when every answer looks equally wrong and suspicious",
                "5 minutes into the exam and already accepting my fate",
                "Study what you love, they said. It'll be fun, they said.",
                "Me looking at the exam paper: well, this is how I die"
            ],
            "studying": [
                "Opening textbook: confidence level 100%. Reading first paragraph: confidence level 0%",
                "My notes look like they were written during an earthquake",
                "Studying with music vs studying in silence: both result in zero studying",
                "My attention span: shorter than a TikTok video",
                "Rereading the same paragraph 10 times and still not processing it"
            ],
            "friendship": [
                "Found the one friend who's as lost as me in class: soulmates",
                "Friendship is sending 'I'm gonna fail' texts at 2 AM and getting 'me too' replies",
                "Calling each other to make sure we're both equally unprepared",
                "Sharing one pizza between four people: the true bond of friendship"
            ],
            "coffee": [
                "Coffee: turning 'I can't human' into 'I can maybe human'",
                "My relationship with coffee: toxic but necessary",
                "Running on coffee, panic, and poor life choices",
                "Spending more on coffee than on groceries: the student way"
            ]
        }
        
        # Try AI first (optional)
        ai_caption = self._try_ai_caption(theme, style)
        if ai_caption:
            return ai_caption
        
        # Use manual database
        import random
        
        theme_lower = theme.lower()
        
        # Direct match
        if theme_lower in CAPTION_DB:
            return random.choice(CAPTION_DB[theme_lower])
        
        # Keyword match
        for key, captions in CAPTION_DB.items():
            if any(word in theme_lower for word in key.split()):
                return random.choice(captions)
        
        # General funny fallbacks
        fallbacks = [
            f"When {theme} hits different than expected",
            f"The {theme} struggle is too real",
            f"{theme.title()}: easier said than done",
            f"My relationship with {theme}: it's complicated",
            f"That {theme} life chose violence today"
        ]
        
        return random.choice(fallbacks)

    def _try_ai_caption(self, theme: str, style: str) -> str:
        """Try AI generation, return None if fails"""
        if not self.model:
            return None
        
        try:
            prompt = f"Create one specific funny meme caption about '{theme}' for college students. Make it relatable and viral. Return only the caption text."
            response = self.model.generate_content(prompt)
            
            if hasattr(response, 'text') and response.text:
                caption = response.text.strip()
                if len(caption) > 20 and "meme" not in caption.lower():
                    return caption
        except:
            pass
        
        return None

    def add_text_to_image(self, image: Image.Image, text: str, position: str = "bottom") -> Image.Image:
        img_copy = image.copy()
        draw = ImageDraw.Draw(img_copy)
        img_width, img_height = img_copy.size

        # ✅ Use your own font file, no system fonts
        font_path = "assets/fonts/Impacted Regular.ttf"

        # Scale font size based on image height (much more consistent)
        font_size = max(40, int(img_height * 0.08))
        font = ImageFont.truetype(font_path, font_size)

        # Split text into lines that fit image width
        words = text.split()
        lines, current_line = [], []
        for word in words:
            test_line = " ".join(current_line + [word])
            bbox = draw.textbbox((0, 0), test_line, font=font)
            if bbox[2] - bbox[0] <= img_width * 0.85:
                current_line.append(word)
            else:
                lines.append(" ".join(current_line))
                current_line = [word]
        if current_line:
            lines.append(" ".join(current_line))

        line_height = font_size * 1.3
        total_height = line_height * len(lines)
        if position == "top":
            y = 20
        elif position == "center":
            y = (img_height - total_height) // 2
        else:
            y = img_height - total_height - 20

        outline = int(font_size * 0.1)  # scalable outline

        # Draw outlined text
        for i, line in enumerate(lines):
            bbox = draw.textbbox((0, 0), line, font=font)
            x = (img_width - (bbox[2] - bbox[0])) // 2
            y_line = y + i * line_height
            for dx in [-outline, 0, outline]:
                for dy in [-outline, 0, outline]:
                    if dx != 0 or dy != 0:
                        draw.text((x + dx, y_line + dy), line, font=font, fill="black")
            draw.text((x, y_line), line, font=font, fill="white")

        return img_copy



def create_sample_templates():
    """Create sample meme templates including custom ones"""
    templates = []
    
    # Your existing gradient templates
    colors = [
        ("#FF6B6B", "#4ECDC4"),  # Red to teal
        ("#845EC2", "#FF9671"),  # Purple to orange
        ("#2E8B57", "#FFD700"),  # Green to gold
        ("#4169E1", "#FF1493"),  # Blue to pink
        ("#8B4513", "#F0E68C"),  # Brown to khaki
    ]

    for i, (color1, color2) in enumerate(colors):
        # Create gradient background
        img = Image.new("RGB", (800, 600), color1)
        draw = ImageDraw.Draw(img)

        # Add gradient effect
        for y in range(600):
            r1, g1, b1 = tuple(int(color1[j:j+2], 16) for j in (1, 3, 5))
            r2, g2, b2 = tuple(int(color2[j:j+2], 16) for j in (1, 3, 5))

            ratio = y / 600
            r = int(r1 * (1 - ratio) + r2 * ratio)
            g = int(g1 * (1 - ratio) + g2 * ratio)
            b = int(b1 * (1 - ratio) + b2 * ratio)

            draw.line([(0, y), (800, y)], fill=(r, g, b))

        # Add decorative elements
        if i % 2 == 0:
            # Add circles
            for x in range(0, 800, 200):
                for y in range(0, 600, 200):
                    draw.ellipse([x-50, y-50, x+50, y+50], outline="white", width=3)
        else:
            # Add rectangles
            for x in range(100, 700, 200):
                for y in range(100, 500, 200):
                    draw.rectangle([x-30, y-30, x+30, y+30], outline="white", width=3)

        templates.append((f"Gradient Template {i+1}", img))

    # ADD YOUR CUSTOM TEMPLATES HERE
    custom_templates = create_custom_templates()
    templates.extend(custom_templates)
    
    return templates

def create_custom_templates():
    """Create your custom meme templates without emojis"""
    custom_templates = []
    
    # TEMPLATE 1: Popular Meme Format - Drake (FIXED)
    drake_template = Image.new("RGB", (800, 600), "white")
    draw = ImageDraw.Draw(drake_template)
    
    # Draw two panels (like Drake meme)
    draw.rectangle([50, 50, 375, 275], outline="black", width=3)  # Top panel
    draw.rectangle([50, 325, 375, 550], outline="black", width=3)  # Bottom panel
    
    # Add labels WITHOUT EMOJIS
    try:
        # Try to use a better font
        font = ImageFont.truetype("arial.ttf", 20)
    except:
        font = ImageFont.load_default()
    
    draw.text((150, 280), "GOOD", fill="black", font=font)
    draw.text((150, 555), "BAD", fill="black", font=font)
    
    custom_templates.append(("Drake Format", drake_template))
    
    # TEMPLATE 2: Distracted Boyfriend (FIXED)
    boyfriend_template = Image.new("RGB", (800, 600), "white")
    draw = ImageDraw.Draw(boyfriend_template)
    
    # Simple representation
    draw.ellipse([200, 200, 300, 300], outline="blue", width=3)  # Boyfriend
    draw.ellipse([400, 150, 450, 200], outline="red", width=3)   # Other girl
    draw.ellipse([100, 250, 150, 300], outline="green", width=3) # Girlfriend
    
    # Arrows
    draw.line([250, 250, 400, 175], fill="red", width=2)
    draw.line([250, 250, 125, 275], fill="green", width=2)
    
    # Labels without emojis
    draw.text([420, 210], "OTHER", fill="red", font=font)
    draw.text([60, 310], "GIRLFRIEND", fill="green", font=font)
    draw.text([210, 310], "BOYFRIEND", fill="blue", font=font)
    
    custom_templates.append(("Distracted Boyfriend", boyfriend_template))
    
    # TEMPLATE 3: Two Buttons (FIXED)
    buttons_template = Image.new("RGB", (800, 600), "white")
    draw = ImageDraw.Draw(buttons_template)
    
    # Two buttons
    draw.rectangle([150, 200, 350, 300], outline="blue", width=3, fill="#e6f3ff")
    draw.rectangle([450, 200, 650, 300], outline="red", width=3, fill="#ffe6e6")
    
    draw.text((200, 230), "BUTTON 1", fill="blue", font=font)
    draw.text((500, 230), "BUTTON 2", fill="red", font=font)
    
    custom_templates.append(("Two Buttons", buttons_template))
    
    # TEMPLATE 4: Brain/Mind Blown (FIXED)
    brain_template = Image.new("RGB", (800, 600), "black")
    draw = ImageDraw.Draw(brain_template)
    
    # Simple brain representation
    for i in range(4):
        y_pos = 150 + i * 100
        draw.ellipse([300, y_pos, 500, y_pos + 80], outline="white", width=2)
    
    # Explosion lines
    for angle in range(0, 360, 30):
        import math
        end_x = 400 + 100 * math.cos(math.radians(angle))
        end_y = 400 + 100 * math.sin(math.radians(angle))
        draw.line([400, 400, end_x, end_y], fill="yellow", width=2)
    
    draw.text((350, 500), "MIND BLOWN", fill="white", font=font)
    
    custom_templates.append(("Brain Explosion", brain_template))
    
    # TEMPLATE 5: Change My Mind
    changemymind = Image.new("RGB", (800, 600), "blue")
    draw = ImageDraw.Draw(changemymind)
    draw.rectangle([100, 100, 700, 500], fill="white", outline="black", width=3)
    draw.text((350, 300), "CHANGE MY MIND", fill="black", font=font)
    custom_templates.append(("Change My Mind", changemymind))
    
    # TEMPLATE 6: Simple Frame
    frame = Image.new("RGB", (800, 600), "#f0f8ff")
    draw = ImageDraw.Draw(frame)
    draw.rectangle([50, 50, 750, 550], outline="#4169e1", width=10)
    draw.text((300, 300), "YOUR MEME HERE", fill="#4169e1", font=font)
    custom_templates.append(("Blue Frame", frame))
    
    return custom_templates

def main():
    # Header
    st.title("🎭 AI Meme Poster Creator")
    st.markdown("Create amazing memes and posters with AI-powered captions!")

    # Initialize meme generator
    meme_gen = MemeGenerator()

    # Sidebar for options
    st.sidebar.header("🎨 Creation Options")

    # Template selection or upload
    option = st.sidebar.radio("Choose Image Source:", ["Use Template", "Upload Image"])

    if option == "Use Template":
        templates = create_sample_templates()
        template_names = [name for name, _ in templates]
        selected_template = st.sidebar.selectbox("Select Template:", template_names)

        # Get selected template image
        selected_image = None
        for name, img in templates:
            if name == selected_template:
                selected_image = img
                break
    else:
        uploaded_file = st.sidebar.file_uploader("Upload Image:", type=['png', 'jpg', 'jpeg'])
        if uploaded_file:
            selected_image = Image.open(uploaded_file)
        else:
            selected_image = None

    if selected_image:
        # Display original image
        col1, col2 = st.columns(2)

        with col1:
            st.subheader("📷 Original Image")
            st.image(selected_image, use_column_width=True)

        # SIDEBAR - CAPTION SECTION
        st.sidebar.header("✍️ Caption Options")

        # Initialize session state values
        if 'caption_source' not in st.session_state:
            st.session_state.caption_source = "Write Custom"
        if 'custom_caption' not in st.session_state:
            st.session_state.custom_caption = ""
        if 'ai_theme' not in st.session_state:
            st.session_state.ai_theme = ""
        if 'ai_style' not in st.session_state:
            st.session_state.ai_style = "funny"
        if 'current_caption' not in st.session_state:
            st.session_state.current_caption = ""

        # Caption source selection
        caption_option = st.sidebar.radio(
            "Caption Source:", 
            ["Write Custom", "Generate AI"],
            key="caption_source"
        )

        if caption_option == "Write Custom":
            # Update custom caption from widget
            st.session_state.custom_caption = st.sidebar.text_area(
                "Enter Caption:", 
                value=st.session_state.custom_caption,
                placeholder="Type your caption here..."
            )
            st.session_state.current_caption = st.session_state.custom_caption
            
        else:  # Generate AI
            # Update AI theme from widget
            st.session_state.ai_theme = st.sidebar.text_input(
                "Meme Theme:", 
                value=st.session_state.ai_theme,
                placeholder="e.g., college life, exams, friendship"
            )
            
            # Get AI style from widget WITHOUT assigning to session state
            ai_style = st.sidebar.selectbox(
                "Caption Style:", 
                ["funny", "motivational", "sarcastic", "wholesome"],
                index=["funny", "motivational", "sarcastic", "wholesome"].index(st.session_state.ai_style)
            )
            # Update session state separately
            st.session_state.ai_style = ai_style

            # Generate AI caption button
            if st.sidebar.button("🤖 Generate Caption", key="generate_ai"):
                if st.session_state.ai_theme:
                    with st.spinner("Generating caption..."):
                        ai_caption = meme_gen.generate_caption(st.session_state.ai_theme, st.session_state.ai_style)
                        st.session_state.current_caption = ai_caption
                        st.session_state.custom_caption = ai_caption  # Also update custom field
                        st.sidebar.success("Caption generated!")
                else:
                    st.sidebar.warning("Please enter a theme first!")

            # Display AI generated caption
            if st.session_state.current_caption and st.session_state.ai_theme:
                st.sidebar.write(f"**AI Caption:** {st.session_state.current_caption}")

        # Text position
        text_position = st.sidebar.selectbox("Text Position:", ["bottom", "top", "center"])

        # CREATE MEME SECTION - SEPARATE AND ALWAYS VISIBLE
        st.sidebar.header("🎭 Create Meme")

        # Show caption status
        if st.session_state.current_caption:
            st.sidebar.success(f"✅ Ready to create!")
            st.sidebar.write(f"**Caption:** {st.session_state.current_caption}")
            button_label = "🎭 Create Meme"
        else:
            st.sidebar.warning("⚠️ Enter a caption first")
            button_label = "🎭 Create Meme (Need Caption)"

        # The main create button - ALWAYS VISIBLE
        create_clicked = st.sidebar.button(button_label, type="primary", use_container_width=True)

        if create_clicked:
            if st.session_state.current_caption:
                with st.spinner("Creating your meme..."):
                    try:
                        generated_image = meme_gen.add_text_to_image(
                            selected_image, 
                            st.session_state.current_caption, 
                            text_position
                        )
                        st.session_state.generated_image = generated_image
                        st.sidebar.success("Meme created! Check the right column →")
                        # Force a rerun to display the image
                        st.rerun()
                    except Exception as e:
                        st.sidebar.error(f"Error creating meme: {str(e)}")
            else:
                st.sidebar.error("❌ Please enter or generate a caption first!")

        # Display generated meme in the RIGHT COLUMN
        if st.session_state.generated_image:
            with col2:
                st.subheader("🎉 Generated Meme")
                st.image(st.session_state.generated_image, use_column_width=True)

                # Download button
                img_buffer = io.BytesIO()
                st.session_state.generated_image.save(img_buffer, format='PNG')
                img_buffer.seek(0)

                st.download_button(
                    label="⬇️ Download Meme",
                    data=img_buffer.getvalue(),
                    file_name="generated_meme.png",
                    mime="image/png"
                )
        else:
            with col2:
                st.info("👆 Generate a meme to see it here!")

    else:
        st.info("👈 Please select a template or upload an image to get started!")

        # Show sample templates
        st.subheader("🎨 Available Templates Preview")
        templates = create_sample_templates()

        cols = st.columns(3)
        for i, (name, img) in enumerate(templates):
            with cols[i % 3]:
                st.image(img, caption=name, use_column_width=True)

    # Footer
    st.markdown("---")
    st.markdown("Built with ❤️ using Streamlit and Gemini AI")

if __name__ == "__main__":
    main()