import streamlit as st
import json
import os
import uuid
import re
import requests
from datetime import datetime
from io import BytesIO
from decimal import Decimal  # Add this import for DynamoDB float handling

# Third-party library imports
import boto3
from PIL import Image
import firebase_admin
from firebase_admin import credentials, auth
import pandas as pd
import streamlit_tags as st_tags
from dotenv import load_dotenv

# Load environment variables from .env file if it exists
load_dotenv()

# Detect if running on mobile
def is_mobile():
    # Try to detect mobile browsers based on User-Agent
    try:
        user_agent = st.get_current_user().user_agent
        return any(device in user_agent.lower() for device in ["android", "iphone", "ipad", "mobile"])
    except:
        # If we can't detect, assume it might be mobile for better experience
        return False

# Auto-expand sidebar on mobile
if is_mobile() and "sidebar_expanded" not in st.session_state:
    st.session_state["sidebar_expanded"] = True
    # Note: This doesn't directly control Streamlit's sidebar, but we'll use this flag

# Custom CSS for better mobile experience
st.markdown("""
<style>
    /* Larger buttons for touch interfaces */
    .stButton>button {
        font-size: 18px !important;
        padding: 12px 16px !important;
        margin: 6px 0 !important;
    }
    
    /* Larger text inputs */
    .stTextInput>div>div>input {
        font-size: 16px !important;
        padding: 10px !important;
    }
    
    /* Improve spacing for mobile */
    .block-container {
        padding-top: 2rem !important;
        padding-bottom: 2rem !important;
    }
    
    /* Make form elements more touch-friendly */
    .stSelectbox, .stNumberInput {
        margin-bottom: 12px !important;
    }
    
    /* Visual cue for sidebar toggle on mobile */
    @media (max-width: 768px) {
        [data-testid="stSidebarNav"] {
            position: relative;
        }
        [data-testid="stSidebarNav"]::after {
            content: "👈 Menu";
            position: absolute;
            right: -60px;
            top: 0;
            background: #f0f2f6;
            padding: 8px;
            border-radius: 4px;
            animation: pulse 2s infinite;
            z-index: 1000;
        }
        @keyframes pulse {
            0% { opacity: 1; }
            50% { opacity: 0.6; }
            100% { opacity: 1; }
        }
    }
</style>
""", unsafe_allow_html=True)

# Load AWS credentials using correct HF Secrets
AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY")
AWS_SECRET_KEY = os.getenv("AWS_SECRET_KEY")
AWS_REGION = os.getenv("AWS_REGION", "us-east-1")
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME", "food-image-crowdsourcing")
DYNAMODB_TABLE = os.getenv("DYNAMODB_TABLE", "image_metadata")
HF_API_TOKEN = os.getenv("HF_API_TOKEN", "")  # For Hugging Face Inference API

# Load Firebase credentials
FIREBASE_CONFIG = json.loads(os.getenv("FIREBASE_CONFIG", "{}"))

# Initialize Firebase Admin SDK (Prevent multiple initialization)
if not firebase_admin._apps:
    try:
        cred = credentials.Certificate(FIREBASE_CONFIG)
        firebase_admin.initialize_app(cred)
    except Exception as e:
        st.error(f"Firebase initialization error: {e}")
        if st.button("Continue in Demo Mode"):
            st.session_state["demo_mode"] = True
        else:
            st.stop()

# Initialize AWS Services (S3 & DynamoDB)
try:
    s3 = boto3.client(
        "s3",
        aws_access_key_id=AWS_ACCESS_KEY,
        aws_secret_access_key=AWS_SECRET_KEY,
        region_name=AWS_REGION
    )

    dynamodb = boto3.resource(
        "dynamodb", 
        region_name=AWS_REGION, 
        aws_access_key_id=AWS_ACCESS_KEY, 
        aws_secret_access_key=AWS_SECRET_KEY,
    )
    metadata_table = dynamodb.Table(DYNAMODB_TABLE)
except Exception as e:
    st.error(f"AWS initialization error: {e}")
    if st.button("Continue in Demo Mode"):
        st.session_state["demo_mode"] = True
    else:
        st.stop()

# Food Intellisense List
FOOD_SUGGESTIONS = [
    "Ajvar", "Angel Wings", "Apple", "Apple Pie", "Apfelstrudel", "Arancini", "Asparagus", "Babka", "Bagel","Baguette", "Baklava",
    "Banana", "Banana Bread", "Banh Mi", "Banitsa", "Barbecue Ribs", "BBQ Chicken", "BBQ Chicken Pizza", "BBQ Ribs", "Bean Buritto",
    "Bear Claw", "Beef Empanadas", "Beef Pho", "Beef Sirloin", "Beef Stroganoff", "Beer", "Beets", "Bell Pepper", "Biryani", "Bistecca alla Fiorentina",
    "Black Beans", "Black Forest Cake", "Black Olives", "Blini", "Borscht", "Bossam", "Brioche", "Broccoli", "Brown Rice", 
    "Bruschetta", "Brussels Sprouts", "Buckwheat", "Buffalo Wings", "Burger", "Burrito", "Butter Chicken", "Cabbage", 
    "Cabbage Rolls", "Calzone", "Cannoli", "Carrot", "Carrot Cake", "Cauliflower", "Cauliflower Soup", "Cevapi", "Ceviche", "Ceviche de Camaron",
    "Challah", "Char Siu", "Cheese Empanadas", "Cheesecake", "Chicken", "Chicken Broth", "Chicken Empanadas", 
    "Chicken Wings", "Chickpeas", "Chiles en Nogada", "Chili Sauce", "Chimichirri Steak", "Chow Mein", 
    "Clams", "Cold Beet Soup", "Corn", "Corn on the Cob", "Coxinha", "Crab Cakes", "Cream Cheese", "Creamy Mushroom Risotto",
    "Creme Brulee", "Creole Gumbo", "Croissant", "Croque Monsieur", "Cucumber", "Cucumber Soup", "Deep-fried", 
    "Dim Sum", "Dolmades", "Doughnuts", "Duck", "Eggplant", "Eggplant Spread", "Eggs", "Enchiladas",
    "Encebollado", "Falafel", "Fanesca", "Fasolada", "Faworki", "Filet Mignon", "Fish", "Fish and Chips",
    "Fish Tacos", "Flatbread", "Flan", "Focaccia", "Four Cheese Pizza", "French Fries", "French Onion Soup",
    "Fresh Fruit", "Fruit Soup", "Garbanzo", "Garlic", "Gazpacho", "Gefilte Fish", "Gibanica", "Ginger Bread",
    "Goat Cheese", "Goulash", "Green Beans", "Green Fried Tomatoes", "Green Onion", "Gyoza", "Gyros", "Hawaiian Pizza",
    "Herbs", "Hoddeok", "Hot and Sour Soup", "Hot Pot", "Hummus", "Hunter's stew", "Ice Cream", "Japchae",
    "Jasmine Rice", "Jollof Rice", "Kabsa", "Kale", "Katsu Curry", "Kavarma", "Kebabs", "Kimchi Fried Rice", "Kisiel",
    "Kremowka", "Kreplach", "Kung Pao Chicken", "Kutia", "Lamb", "Lamb Chops", "Lasagna", "Layered Potato Casserole",
    "Lemon", "Lemon Pie", "Lentil Soup", "Lettuce", "Llapingachos", "Lobster", "Mac and Cheese", "Macarons", "Mahi Mahi", 
    "Mansaf", "Mapo Tofu", "Margherita Pizza", "Marinated", "Marzipan", "Matzo Ball Soup", "Mazurek", "Meat Lover's Pizza",
    "Meat Patties", "Meatloaf", "Miso Soup", "Mixed Salad", "Mixed Vegetables", "Mooncake", "Moussaka", "Mozarella", "Mushroom Pizza", "Mushroom Soup",
    "Mushrooms", "Napoleon Cake", "Neapolitan Pizza", "New York Strip Steak", "Nougat Candies", "Onion Rings", "Onion",
    "Osso Buco", "Oysters", "Pad Thai", "Paella", "Panna Cotta", "Pasta", "Pasta Carbonara", "Pavlova", 
    "Peas", "Pecan Pie", "Peking Duck", "Pelmeni", "Pepperoni Pizza", "Pierogi", "Pineapple", "Pita Bread",
    "Pizza", "Pljeskavica", "Pork Chops", "Pork Knuckle", "Portobello Mushrooms", "Potato pancakes", "Potato Salad",
    "Poutine", "Poppy Seed Roll", "Pudding", "Pulled Pork", "Pumpkin", "Pumpkin Pie", "Radish", "Quesadillas", "Quiche", "Ramen", "Ratatouille",
    "Ravioli", "Red Pepper", "Ribeye Steak", "Ribolita", "Rich Stew", "Risotto alla Milanese", "Roll (Multi-grain)", "Roll (Multigrain)", 
    "Roll (Poppyseed)","Roll (Rye)", "Roll (Sesame)", "Roll (Sourdough)", "Roll (Wheat)", "Roll (White)", "Rugelach", "Rye Bread",
    "Sachertorte", "Saffron Rice", "Salad", "Salmon", "Sarma", "Sausage", "Sauerkraut", "Seafood Pasta", 
    "Seco de Chivo", "Shashlik", "Shashuka", "Shawarma", "Shepherd's Pie", "Shopska Salad", "Shrimp", "Shrimp Skewers",
    "Soft Egg Noodles", "Sopes", "Soup Dumplings", "Sour-Dough Bread", "Sour Rye Soup", "Souvlaki", "Spaghetti Carbonara", "Spinach", "Sponge Cake",
    "Spring Salad", "Spring Rolls", "Stuffed Cabbage", "Stuffed Grape Leaves", "Stuffed Mushrooms", "Stuffed Pepper", "Supreme Pizza", "Sushi", 
    "Swwet and Sour Pork", "Sweet Potato", "Swordfish Steak", "Szarlotka", "T-bone Steak", "Tacos", "Tamales", "Tandoori Chicken", "Teriyaki", "Tarator",
    "Texas Style Brisket", "Tilapia", "Tiramisu", "Toast", "Tomato", "Tomato Soup", "Tostada", "Tteokbokki", "Tuna Steak",
    "Tzatziki", "Uszka", "Vareniki", "Veal", "Veggie Fries", "Veggie Pizza", "Wheat Bread", "White Bean Soup", "White Pizza", 
    "Wiener Schnitzel", "Wild Mushroom Pasta", "Wine (Red)", "Wine (White)", "Wonton Soup", "Xiaolongbao", "Zeppelins", "Zucchini"
]  # Alphabetically sorted list of diverse cuisines

# Unit options for food weight/volume
UNIT_OPTIONS = ["grams", "ounce(s)", "teaspoon(s)", "tablespoon(s)", "cup(s)", "slice(s)", "piece(s)"]

# Cooking methods
COOKING_METHODS = [
    "Unknown", "Baked", "Boiled", "Braised", "Breaded and fried", "Broiled", "Creamy", "Deep-fried", "Dried", 
    "Fried", "Grilled", "Grilled minced", "Marinated", "Microwaved", "Pan-seared", "Poached", "Raw", 
    "Roasted", "Sautéed", "Slow-cooked", "Smoked", "Steamed", "Stewed", "Stir-fried", "Takeout/Restaurant"
]

# Helper functions
def resize_image(image, max_size=512, quality=85):
    """
    Resize image while preserving aspect ratio and reducing file size
    
    Args:
        image: PIL Image object
        max_size: Maximum dimension (width or height)
        quality: JPEG quality (0-100)
        
    Returns:
        Resized PIL Image
    """
    # Calculate new dimensions
    width, height = image.width, image.height
    
    # Only resize if the image is larger than max_size
    if width > max_size or height > max_size:
        if width > height:
            new_width = max_size
            new_height = int(height * (max_size / width))
        else:
            new_height = max_size
            new_width = int(width * (max_size / height))
        
        # Resize the image
        resized_img = image.resize((new_width, new_height), Image.LANCZOS)
    else:
        # If image is already smaller than max_size, return a copy to avoid modifying original
        resized_img = image.copy()
    
    # Convert to RGB if image has alpha channel (for JPEG conversion)
    if resized_img.mode == 'RGBA':
        resized_img = resized_img.convert('RGB')
    
    # Compress the image
    buffer = BytesIO()
    resized_img.save(buffer, format="JPEG", quality=quality, optimize=True)
    buffer.seek(0)
    
    # Return the compressed image
    return Image.open(buffer)

def get_image_size_kb(image):
    """Get image file size in KB"""
    buffer = BytesIO()
    image.save(buffer, format="JPEG")
    size_bytes = buffer.tell()
    return size_bytes / 1024  # Convert to KB

def upload_to_s3(image, user_id, folder="", force_quality=None):
    """
    Upload image to S3 bucket and return the S3 path
    
    Args:
        image: PIL Image object
        user_id: User ID for folder structure
        folder: Subfolder to store the image in (e.g., "raw-uploads" or "processed-512x512")
        force_quality: Override default quality settings if specified
    """
    if st.session_state.get("demo_mode", False):
        return f"demo/{user_id}/demo_image.jpg"
    
    try:
        # Generate a unique ID for the image
        image_id = str(uuid.uuid4())
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        
        # Create the S3 path with the appropriate folder structure
        if folder:
            s3_path = f"{folder}/{user_id}/{timestamp}_{image_id}.jpg"
        else:
            s3_path = f"{user_id}/{timestamp}_{image_id}.jpg"
        
        # Convert PIL image to bytes
        buffer = BytesIO()
        
        # Set quality based on folder or forced value
        if force_quality is not None:
            quality = force_quality
        else:
            # Higher quality for raw uploads, compressed for processed
            quality = 95 if folder == "raw-uploads" else 85
        
        # Don't compress the image again if it's already been through resize_image
        # Just save with the appropriate quality
        image.save(buffer, format="JPEG", quality=quality, optimize=True)
        buffer.seek(0)
        
        # Upload to S3
        s3.upload_fileobj(buffer, S3_BUCKET_NAME, s3_path)
        return s3_path
    except Exception as e:
        st.error(f"Failed to upload image: {e}")
        return None

def transcribe_audio(audio_bytes):
    """Transcribe audio using Hugging Face's Whisper model via Inference API"""
    try:
        # Convert audio bytes to file-like object
        audio_file = BytesIO(audio_bytes)
        
        # Free Hugging Face Inference API endpoint for Whisper Tiny model
        API_URL = "https://api-inference.huggingface.co/models/openai/whisper-tiny"
        
        headers = {}
        if HF_API_TOKEN:
            headers["Authorization"] = f"Bearer {HF_API_TOKEN}"
        
        # Make request to the free HF API
        response = requests.post(
            API_URL,
            headers=headers,
            data=audio_file
        )
        
        if response.status_code == 200:
            result = response.json()
            # Extract text from response
            transcript = result.get("text", "")
            return transcript
        else:
            # Fallback for rate limiting or errors
            st.warning("Could not transcribe audio. Please try typing instead.")
            return ""
            
    except Exception as e:
        st.error(f"Transcription error: {e}")
        return ""

def parse_food_annotation(transcript, focus_fields=None):
    """
    Parse the transcribed text to extract food details
    Simple rule-based parsing for common patterns
    Optional focus_fields parameter to prioritize specific fields
    """
    # Default values
    parsed_data = {
        "food_name": "",
        "portion_size": None,
        "portion_unit": "",
        "cooking_method": "Unknown",
        "ingredients": []
    }
    
    # Try to extract food name
    # Start with items from our suggestion list
    for food in FOOD_SUGGESTIONS:
        if food.lower() in transcript.lower():
            parsed_data["food_name"] = food
            break
    
    # If no match, use the first few words as the food name
    if not parsed_data["food_name"]:
        words = transcript.split()
        if words:
            # Use first 3 words or less as food name
            parsed_data["food_name"] = " ".join(words[:min(3, len(words))])
    
    # Try to extract portion size and unit
    # Look for patterns like "100 grams" or "2 slices"
    size_match = re.search(r'(\d+(?:\.\d+)?)\s*(grams?|ounces?|cups?|pieces?|slices?)', transcript.lower())
    if size_match:
        try:
            parsed_data["portion_size"] = float(size_match.group(1))
            # Map to our standard units
            unit_text = size_match.group(2).rstrip('s')  # Remove plural 's'
            if unit_text == "gram":
                parsed_data["portion_unit"] = "grams"
            elif unit_text == "ounce":
                parsed_data["portion_unit"] = "ounce(s)"
            elif unit_text == "cup":
                parsed_data["portion_unit"] = "cup(s)"
            elif unit_text == "slice":
                parsed_data["portion_unit"] = "slice(s)"
            elif unit_text == "piece":
                parsed_data["portion_unit"] = "piece(s)"
        except:
            pass
    
    # Try to extract cooking method
    for method in COOKING_METHODS:
        if method.lower() in transcript.lower():
            parsed_data["cooking_method"] = method
            break
    
    # Simple ingredient extraction
    common_ingredients = ["cheese", "tomato", "lettuce", "onion", "beef", "chicken", "salt", "pepper"]
    found_ingredients = []
    for ingredient in common_ingredients:
        if ingredient.lower() in transcript.lower():
            found_ingredients.append(ingredient.capitalize())
    
    if found_ingredients:
        parsed_data["ingredients"] = found_ingredients
    
    # If focus_fields is provided, prioritize extracting those fields
    if focus_fields:
        # More targeted extraction methods for specific fields
        if "food_name" in focus_fields:
            # More aggressive food name extraction
            # e.g., assume the entire transcript might be just the food name
            if not parsed_data["food_name"]:
                parsed_data["food_name"] = transcript.strip()
        
        if "portion_size" in focus_fields or "portion_unit" in focus_fields:
            # More aggressive portion extraction
            # e.g., assume numbers are portion sizes even without units
            if not parsed_data["portion_size"]:
                number_match = re.search(r'(\d+(?:\.\d+)?)', transcript)
                if number_match:
                    parsed_data["portion_size"] = float(number_match.group(1))
                    parsed_data["portion_unit"] = "piece(s)"  # Default unit
    
    return parsed_data

def save_metadata(user_id, s3_path, food_name, portion_size, portion_unit, cooking_method, ingredients, tokens_awarded):
    """Save metadata to DynamoDB"""
    if st.session_state.get("demo_mode", False):
        st.success("Demo mode: Metadata would be saved to DynamoDB")
        return True
    
    try:
        # Generate a unique ID for the database entry
        image_id = str(uuid.uuid4())
        timestamp = datetime.now().isoformat()
        
        # Ensure portion_size is a Decimal (DynamoDB doesn't support float)
        if not isinstance(portion_size, Decimal):
            portion_size = Decimal(str(portion_size))
        
        # Create item for DynamoDB
        item = {
            'image_id': image_id,
            'user_id': user_id,
            'upload_timestamp': timestamp,
            'food_name': food_name,
            'portion_size': portion_size,  # Decimal type
            'portion_unit': portion_unit,
            'cooking_method': cooking_method,
            'ingredients': ingredients,
            's3_path': s3_path,
            'tokens_awarded': tokens_awarded
        }
        
        # Save to DynamoDB
        metadata_table.put_item(Item=item)
        return True
    except Exception as e:
        st.error(f"Failed to save metadata: {e}")
        return False

def calculate_tokens(image_quality, has_metadata, is_unique_category):
    """Calculate tokens based on various factors"""
    tokens = 1  # Base token for upload
    
    if image_quality == "high":
        tokens += 1
    
    if has_metadata:
        tokens += 1
    
    if is_unique_category:
        tokens += 1
    
    return tokens

# Initialize session state for first-time users
if "tokens" not in st.session_state:
    st.session_state["tokens"] = 0

if "uploads_count" not in st.session_state:
    st.session_state["uploads_count"] = 0

# Initialize food items list for storing multiple annotations
if "food_items" not in st.session_state:
    st.session_state["food_items"] = []

# Initialize form input state variables
if "custom_food_name" not in st.session_state:
    st.session_state["custom_food_name"] = ""
if "form_key" not in st.session_state:
    st.session_state["form_key"] = 0  # Add a form key to force re-rendering

# Track partial annotation state for audio recording
if "partial_annotation" not in st.session_state:
    st.session_state["partial_annotation"] = {
        "food_name": "",
        "portion_size": None,
        "portion_unit": "",
        "cooking_method": "",
        "ingredients": []
    }
    
if "missing_fields" not in st.session_state:
    st.session_state["missing_fields"] = []

def reset_form_fields():
    """Reset all form fields after adding an item by incrementing the form key"""
    # Reset custom food name
    st.session_state["custom_food_name"] = ""
    # Increment the form key to force re-rendering with default values
    st.session_state["form_key"] = st.session_state.get("form_key", 0) + 1

def add_food_item(food_name, portion_size, portion_unit, cooking_method, ingredients):
    """Add a food item to the session state"""
    # Set cooking method to "Unknown" if empty
    if not cooking_method:
        cooking_method = "Unknown"
        
    if food_name and portion_size and portion_unit:  # Cooking method no longer required
        # Add the food item to the session state
        st.session_state["food_items"].append({
            "food_name": food_name,
            "portion_size": portion_size,
            "portion_unit": portion_unit,
            "cooking_method": cooking_method,
            "ingredients": ingredients
        })
        st.success(f"✅ Added {food_name} to your submission")
        reset_form_fields()  # Reset form by incrementing key
        return True
    else:
        st.error("❌ Please fill in all required fields")
        return False

# Main App UI
def main():
    # Check if we should display the mobile welcome dialog
    if is_mobile() and "mobile_welcome_shown" not in st.session_state:
        st.session_state["mobile_welcome_shown"] = True
        # Show welcome message for first-time mobile users
        st.info("👋 Welcome to the Food Image Crowdsourcing App! Tap the menu icon (≡) in the top-right corner to login.")

    # Improved authentication UI for mobile
    if is_mobile():
        # Show prominent login button if not logged in
        if "user_id" not in st.session_state and not st.session_state.get("demo_mode", False):
            st.title("🍽️ Food Image Crowdsourcing")
            auth_container = st.container()
            auth_container.warning("⚠️ Please login to continue")
            
            # Big prominent login buttons
            login_col1, login_col2 = st.columns(2)
            with login_col1:
                if st.button("📱 LOGIN", use_container_width=True, type="primary"):
                    st.session_state["sidebar_expanded"] = True
                    st.rerun()
            with login_col2:
                if st.button("✍️ SIGN UP", use_container_width=True):
                    st.session_state["sidebar_expanded"] = True
                    st.rerun()
            
            st.markdown("### 🍕 Help us collect food images!")
            st.markdown("Take pictures of your meals, label them, and earn tokens!")
            
            # Add links to guidelines and terms
            st.markdown("### 📚 Learn More")
            with st.expander("📋 How It Works"):
                try:
                    with open("PARTICIPATION_GUIDELINES.md", "r") as f:
                        guidelines = f.read()
                    st.markdown(guidelines, unsafe_allow_html=True)
                except Exception as e:
                    st.error(f"Could not load guidelines: {e}")
            
            with st.expander("🪙 Earn Tokens"):
                try:
                    with open("TOKEN_REWARDS.md", "r") as f:
                        rewards = f.read()
                    st.markdown(rewards, unsafe_allow_html=True)
                except Exception as e:
                    st.error(f"Could not load rewards information: {e}")
            
            st.stop()

    # Streamlit Layout - Authentication Section in Sidebar
    st.sidebar.title("🔑 User Authentication")
    auth_option = st.sidebar.radio("Select an option", ["Login", "Sign Up", "Logout"])

    if auth_option == "Sign Up":
        email = st.sidebar.text_input("Email")
        password = st.sidebar.text_input("Password", type="password")
        if st.sidebar.button("Sign Up"):
            try:
                if st.session_state.get("demo_mode", False):
                    st.sidebar.success("✅ Demo mode: User created successfully! Please log in.")
                else:
                    user = auth.create_user(email=email, password=password)
                    st.sidebar.success("✅ User created successfully! Please log in.")
                    # Show continue button after signup
                    if st.sidebar.button("▶️ Continue to Login"):
                        st.rerun()
            except Exception as e:
                st.sidebar.error(f"Error: {e}")

    if auth_option == "Login":
        email = st.sidebar.text_input("Email")
        password = st.sidebar.text_input("Password", type="password")
        if st.sidebar.button("Login"):
            try:
                if st.session_state.get("demo_mode", False):
                    st.session_state["user_id"] = "demo_user_123"
                    st.session_state["tokens"] = 0  # Initialize token count
                    st.sidebar.success("✅ Demo mode: Logged in successfully!")
                    # Show continue button after login
                    if st.sidebar.button("▶️ Continue to App"):
                        st.rerun()
                else:
                    user = auth.get_user_by_email(email)
                    st.session_state["user_id"] = user.uid
                    st.session_state["tokens"] = 0  # Initialize token count
                    st.sidebar.success("✅ Logged in successfully!")
                    # Show continue button after login
                    if st.sidebar.button("▶️ Continue to App"):
                        st.rerun()
            except Exception as e:
                st.sidebar.error(f"Login failed: {e}")

    if auth_option == "Logout" and "user_id" in st.session_state:
        del st.session_state["user_id"]
        st.sidebar.success("✅ Logged out successfully!")

    # Ensure user is logged in before uploading
    if "user_id" not in st.session_state and not st.session_state.get("demo_mode", False):
        st.warning("⚠️ Please log in to upload images.")
        
        # Add links to guidelines and terms
        st.markdown("### 📚 While You're Here")
        st.markdown("Take a moment to read our guidelines and token system:")
        
        # Use expanders instead of columns for better document display
        with st.expander("📋 Participation Guidelines"):
            try:
                with open("PARTICIPATION_GUIDELINES.md", "r") as f:
                    guidelines = f.read()
                st.markdown(guidelines, unsafe_allow_html=True)
            except Exception as e:
                st.error(f"Could not load guidelines: {e}")
        
        with st.expander("🪙 Token Rewards System"):
            try:
                with open("TOKEN_REWARDS.md", "r") as f:
                    rewards = f.read()
                st.markdown(rewards, unsafe_allow_html=True)
            except Exception as e:
                st.error(f"Could not load rewards information: {e}")
        
        with st.expander("📜 Terms of Service"):
            try:
                with open("TERMS_OF_SERVICE.md", "r") as f:
                    terms = f.read()
                st.markdown(terms, unsafe_allow_html=True)
            except Exception as e:
                st.error(f"Could not load terms: {e}")
        
        st.stop()

    # Streamlit Layout - Main App
    st.title("🍽️ Food Image Review & Annotation")

    # Compliance & Disclaimer Section
    with st.expander("📜 Terms & Conditions", expanded=False):
        st.markdown("### **Terms & Conditions**")
        st.write(
            "By uploading an image, you agree to transfer full copyright to the research team for AI training purposes."
            " You are responsible for ensuring you own the image and it does not violate any copyright laws."
            " We do not guarantee when tokens will be redeemable. Keep track of your user ID.")
        terms_accepted = st.checkbox("I agree to the terms and conditions", key="terms_accepted")
        if not terms_accepted:
            st.warning("⚠️ You must agree to the terms before proceeding.")
            st.stop()

    # Mobile-friendly workflow indicator
    if is_mobile():
        # Show a progress indicator at the top
        st.markdown("### 📱 Mobile Workflow")
        workflow_steps = ["📷 Upload Image", "🔍 Review Image", "🏷️ Add Food Details", "📤 Submit"]
        
        # Determine current step
        current_step = 0
        if "original_image" in st.session_state:
            current_step = 1
            if st.session_state["food_items"]:
                current_step = 2
        
        # Display steps with highlight on current
        step_cols = st.columns(len(workflow_steps))
        for i, (col, step) in enumerate(zip(step_cols, workflow_steps)):
            if i == current_step:
                col.markdown(f"**{step}** ✓")
            else:
                col.markdown(f"{step}")
                
        st.markdown("---")

    # Upload Image - Larger and more prominent on mobile
    if is_mobile():
        st.markdown("### 📷 Take or Upload a Food Photo")
        st.info("Take a picture of your meal or upload an existing photo")
    
    uploaded_file = st.file_uploader("Upload an image of your food", type=["jpg", "png", "jpeg"])
    if uploaded_file:
        original_img = Image.open(uploaded_file)
        st.session_state["original_image"] = original_img

    # If an image has been uploaded, process and display it
    if "original_image" in st.session_state:
        original_img = st.session_state["original_image"]
        
        # Process the image - resize and compress with more visible difference
        processed_img = resize_image(original_img, max_size=512, quality=85)
        st.session_state["processed_image"] = processed_img
        
        # Calculate file sizes
        original_size = get_image_size_kb(original_img)
        processed_size = get_image_size_kb(processed_img)
        size_reduction = ((original_size - processed_size) / original_size) * 100 if original_size > 0 else 0

        # On mobile, stack images vertically instead of side by side
        if is_mobile():
            st.markdown("### 🔍 Review Your Image")
            
            # Original image
            st.subheader("📷 Original Image")
            st.markdown(f"<div style='border:2px solid red;padding:5px;'>", unsafe_allow_html=True)
            st.image(original_img, caption=f"Original ({original_img.width}x{original_img.height} px, {original_size:.1f} KB)", use_container_width=True)
            st.markdown("</div>", unsafe_allow_html=True)
            
            # Processed image
            st.subheader("🖼️ Processed Image")
            st.markdown(f"<div style='border:2px solid green;padding:5px;'>", unsafe_allow_html=True)
            st.image(processed_img, caption=f"Processed ({processed_img.width}x{processed_img.height} px, {processed_size:.1f} KB)", use_container_width=True)
            st.markdown("</div>", unsafe_allow_html=True)
        else:
            # Desktop layout (side by side)
            col1, col2 = st.columns(2)
            with col1:
                st.subheader("📷 Original Image")
                st.markdown(f"<div style='border:2px solid red;padding:5px;'>", unsafe_allow_html=True)
                st.image(original_img, caption=f"Original ({original_img.width}x{original_img.height} px, {original_size:.1f} KB)", use_container_width=True)
                st.markdown("</div>", unsafe_allow_html=True)
            with col2:
                st.subheader("🖼️ Processed Image")
                st.markdown(f"<div style='border:2px solid green;padding:5px;'>", unsafe_allow_html=True)
                st.image(processed_img, caption=f"Processed ({processed_img.width}x{processed_img.height} px, {processed_size:.1f} KB)", use_container_width=True)
                st.markdown("</div>", unsafe_allow_html=True)
        
        # Show size reduction
        if size_reduction > 5:  # Only show if there's a meaningful reduction
            st.success(f"✅ Image size reduced by {size_reduction:.1f}% for faster uploads and processing")

        # Display existing food annotations if any
        if st.session_state["food_items"]:
            st.subheader("📋 Added Food Items")
            for i, item in enumerate(st.session_state["food_items"]):
                with st.expander(f"🍽️ {item['food_name']} ({item['portion_size']} {item['portion_unit']})"):
                    st.write(f"**Cooking Method:** {item['cooking_method']}")
                    st.write(f"**Ingredients:** {', '.join(item['ingredients'])}")
                    if st.button(f"Remove Item #{i+1}", key=f"remove_{i}"):
                        st.session_state["food_items"].pop(i)
                        st.rerun()

        # Food metadata form
        st.subheader("�� Add Food Details")

        # Use Streamlit form to capture Enter key and provide a better UX
        # Use a dynamic key based on form_key to force re-rendering with default values
        form_key = st.session_state.get("form_key", 0)
        with st.form(key=f"food_item_form_{form_key}"):
            food_selection = st.selectbox("Food Name", options=[""] + FOOD_SUGGESTIONS, index=0)
            
            # Only show custom food name if the dropdown is empty
            custom_food_name = ""
            if food_selection == "":
                custom_food_name = st.text_input("Or enter a custom food name", 
                                              value=st.session_state["custom_food_name"])
            
            # Determine the actual food name to use
            food_name = food_selection if food_selection else custom_food_name
            
            col1, col2 = st.columns(2)
            with col1:
                portion_size = st.number_input("Portion Size", 
                                              min_value=0.1, 
                                              step=0.1, 
                                              format="%.2f",
                                              value=0.1)  # Always use default values
            with col2:
                portion_unit = st.selectbox("Unit", 
                                           options=UNIT_OPTIONS,
                                           index=0)  # Always use default values
            
            # Set Cooking Method with "Unknown" as the default (index 0)
            cooking_method = st.selectbox("Cooking Method (optional)", 
                                         options=COOKING_METHODS, 
                                         index=0)  # Always use default values
            
            ingredients = st_tags.st_tags(
                label="Main Ingredients (Add up to 5)",
                text="Press enter to add",
                value=[],
                suggestions=["Salt", "Pepper", "Olive Oil", "Butter", "Garlic", "Onion", "Tomato"],
                maxtags=5
            )
            
            # Submit button inside the form
            submitted = st.form_submit_button(label="➕ Add This Food Item")
            if submitted:
                if add_food_item(food_name, portion_size, portion_unit, cooking_method, ingredients):
                    # Store the custom food name if needed for future use
                    if custom_food_name:
                        st.session_state["custom_food_name"] = custom_food_name
                    # Don't call reset_form_fields() here, it's already called in add_food_item
                    st.rerun()
        
        # Make submit button more prominent
        st.markdown("---")
        
        # More prominent submit button with instructions
        st.markdown("### 📤 Submit Your Food Annotations")
        st.info("⚠️ After adding all your food items, click the button below to save your submission and earn tokens.")
        
        # Create a larger, more visible submit button
        submit_col1, submit_col2, submit_col3 = st.columns([1, 2, 1])
        with submit_col2:
            if st.button("📤 SUBMIT ALL FOOD ITEMS", 
                        disabled=len(st.session_state["food_items"]) == 0,
                        use_container_width=True,
                        type="primary"):
                if not st.session_state["food_items"]:
                    st.error("❌ Please add at least one food item before submitting")
                else:
                    with st.spinner("Processing your submission..."):
                        all_saved = True
                        total_tokens = 0
                        
                        # Determine image quality (simplified version)
                        image_quality = "high" if original_img.width >= 1000 and original_img.height >= 1000 else "standard"
                        
                        # Get original image file size for comparison
                        original_size = get_image_size_kb(original_img)
                        
                        # Ensure we have a properly processed image with the right settings
                        # Force resize and compression with settings that guarantee size reduction
                        processed_img = resize_image(original_img, max_size=512, quality=85)
                        processed_size = get_image_size_kb(processed_img)
                        
                        # If the processed image isn't smaller enough, reduce quality further
                        if processed_size > original_size * 0.8:  # Ensure at least 20% reduction
                            processed_img = resize_image(original_img, max_size=512, quality=70)
                            processed_size = get_image_size_kb(processed_img)
                            
                            # If still not small enough, try more aggressive compression
                            if processed_size > original_size * 0.8:
                                processed_img = resize_image(original_img, max_size=480, quality=60)
                        
                        # Upload original to raw-uploads folder
                        raw_s3_path = upload_to_s3(original_img, st.session_state["user_id"], 
                                                  folder="raw-uploads", force_quality=95)
                        
                        # Upload only one processed image to processed-512x512 folder
                        processed_s3_path = upload_to_s3(processed_img, st.session_state["user_id"], 
                                                        folder="processed-512x512", force_quality=85)
                        
                        if raw_s3_path and processed_s3_path:
                            # Save each food item with the processed image path
                            for food_item in st.session_state["food_items"]:
                                # Check if metadata is complete
                                has_metadata = True  # Already validated
                                
                                # Check if the food is in a unique category (simplified)
                                is_unique_category = food_item["food_name"] not in ["Pizza", "Burger", "Pasta", "Salad"]
                                
                                # Calculate tokens for this item
                                tokens_awarded = calculate_tokens(image_quality, has_metadata, is_unique_category)
                                total_tokens += tokens_awarded
                                
                                # Convert float to Decimal for DynamoDB
                                portion_size_decimal = Decimal(str(food_item["portion_size"]))
                                
                                # Save metadata to DynamoDB with processed image path
                                success = save_metadata(
                                    st.session_state["user_id"],
                                    processed_s3_path,  # Use the processed image path
                                    food_item["food_name"],
                                    portion_size_decimal,  # Use Decimal type
                                    food_item["portion_unit"],
                                    food_item["cooking_method"],
                                    food_item["ingredients"],
                                    tokens_awarded
                                )
                                
                                if not success:
                                    all_saved = False
                                    break
                            
                            if all_saved:
                                st.session_state["tokens"] += total_tokens
                                st.session_state["uploads_count"] += 1
                                st.success(f"✅ All food items uploaded successfully! You earned {total_tokens} tokens.")
                                
                                # Clear the form and image for a new submission
                                st.session_state.pop("original_image", None)
                                st.session_state.pop("processed_image", None)
                                st.session_state["food_items"] = []
                                st.rerun()
                            else:
                                st.error("Failed to save some items. Please try again.")
                        else:
                            st.error("Failed to upload images. Please try again.")

# Display earned tokens
st.sidebar.markdown("---")
st.sidebar.subheader("🏆 Your Statistics")
st.sidebar.info(f"🪙 Total Tokens: {st.session_state['tokens']}")
st.sidebar.info(f"📸 Total Uploads: {st.session_state.get('uploads_count', 0)}")

# Help and Documentation Links
st.sidebar.markdown("---")
st.sidebar.subheader("📚 Resources")
if st.sidebar.button("Participation Guidelines"):
    with open("PARTICIPATION_GUIDELINES.md", "r") as f:
        guidelines = f.read()
    st.sidebar.markdown(guidelines)

if st.sidebar.button("Token Rewards System"):
    with open("TOKEN_REWARDS.md", "r") as f:
        rewards = f.read()
    st.sidebar.markdown(rewards)

if st.sidebar.button("Terms of Service"):
    with open("TERMS_OF_SERVICE.md", "r") as f:
        terms = f.read()
    st.sidebar.markdown(terms)

if __name__ == "__main__":
    main()