import streamlit as st
import base64
from huggingface_hub import upload_file
from openai import OpenAI
from datetime import datetime
import json
import re

# Page configuration - MUST be first Streamlit command
st.set_page_config(
    page_title="Gen AI - Menu Visual Generator",
    layout="wide",
    page_icon="🧪"
)

# API Key Input Section
st.markdown("### 🔑 OpenAI API Configuration")
api_key_input = st.text_input(
    "Enter your OpenAI API Key",
    type="password",
    placeholder="sk-...",
    help="Get your API key from https://platform.openai.com/api-keys"
)

# Initialize OpenAI client with UI input
client = None
if api_key_input:
    try:
        client = OpenAI(api_key=api_key_input)
        st.success("✅ API key provided successfully!")
    except Exception as e:
        st.error(f"❌ Invalid API key: {str(e)}")
        st.stop()

def create_menu_text_from_image(image_file):
    """Extract text from image"""
    try:
        image_bytes = image_file.read()
        base64_image = base64.b64encode(image_bytes).decode('utf-8')

        # Reset file pointer
        image_file.seek(0)

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": """Please extract all menu items from this image and organize them in a structured format. For each item, extract:
                        1. Item name
                        2. Description (if available)
                        3. Price
                        4. Category (appetizers, mains, desserts, bevrages, etc.)
                        
                        Format your response as a JSON structure like this:
                        {
                            "restaurant_image": "Restaurant Image (if visible),
                            "cusine_type": "Indian/Italian/Mexico/British/etc. (best guess)",
                            "categories": {
                                "Appetizers": [
                                    { "name": "Item Name", "description": "Description (if available)", "price": "$x.xx" }
                                ],
                                "Main Cources": [
                                    { "name": "Item Name", "description": "Description (if available)", "price": "$x.xx" }
                                ],
                            }
                        }
                        
                        If you can't determine the category, group similar items together. Be as accurate as possible with the text extraction."""
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"
                        }
                    }
                ]
            }],
            max_tokens=2000,
        )

        menu_text = response.choices[0].message.content

        # Clean the response and extract JSON
        json_start = menu_text.find("{")
        json_end = menu_text.rfind("}") + 1
        if json_start != -1 and json_end != -1:
            menu_data = json.loads(menu_text[json_start:json_end])
            return menu_data
        else:
            return { "error": "Unable to extract menu data from image", "raw_text": menu_text }

    except Exception as e:
        return { "error": f"Error extracting data from image: {str(e)}" }


def get_nutrinutional_info(dish_name, cuisine_type):
    """Get nutrient information from Nutritionix API"""
    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": f"""For the {cuisine_type} dish '{dish_name}', provide estimated nutritional information in JSON format:
                {{
                    "calories": 350,
                    "protein": 10g,
                    "carbs": 30g,
                    "fat": 15g,
                    "fiber": 5g
                }}
                
                Base your estimate on typical restaurant portions and ingredients for this type of dish. Be realistic with the values."""
            }],
            max_tokens=200,
        )
        nutrient_info = response.choices[0].message.content

        # Clean the response and extract JSON
        json_start = menu_text.find("{")
        json_end = menu_text.rfind("}") + 1
        if json_start != -1 and json_end != -1:
            menu_data = json.loads(menu_text[json_start:json_end])
            return menu_data
        else:
            return { "calories": 'N/A', "protein": 'N/A',  "carbs": 'N/A', "fat": 'N/A', "fiber": 'N/A' }

    except Exception as e:
        return { "calories": 'N/A', "protein": 'N/A',  "carbs": 'N/A', "fat": 'N/A', "fiber": 'N/A' }

def generate_food_image(dish_name, description, cuisine_type, image_style):
    """Generate food image with enhanced prompting"""

    # style configurations
    style_config = {
        "Professional Food Photography": {
            "lighting": "professional studio lighting with soft shadows",
            "background": "clean white or neutral background",
            "composition": "centered plating with professional garnish",
            "quality": "restaurant-quality presentation, high-end food photography"
        },
        "Rustic Homestyle": {
            "lighting": "warm, natural lighting",
            "background": "rusted wooden table or textured surface",
            "composition": "casual, homestyle presentation",
            "quality": "comforting, homemade appearance with natrual styling"
        },
        "Modern Minimalstic": {
            "lighting": "clean, bright lighting",
            "background": "minimalistic white or light gray background",
            "composition": "artistic plating with negative space",
            "quality": "contemporary, Instagram-worthy presentation"
        },
        "Vibrant Colorful": {
            "lighting": "brigt, vibrant lighting that enhances color",
            "background": "colorful or complementary background",
            "composition": "dynamic, eye-catching, presentation",
            "quality": "bold, appetizing colors that pop"
        }
    }

    selected_style = style_config.get(image_style, style_config['Professional Food Photography'])

    # Create detailed food photography prompt
    food_prompt = f"""
Professional food photography of {dish_name} ({cuisine_type} cuisine).
    
Dish details: {description}
    
Photography specifications:
- {selected_style['lighting']}
- {selected_style['background']}
- {selected_style['composition']}
- {selected_style['quality']}

Technical Requirements:
- High resolution, commercial food photography quality
- Appetizing and mouth-watering presentation
- Perfect focus and sharp details
- Colors that enhance appetite appeal
- Professional plating and garnish
- Shot from optimal angle to showcase the dish
- No text or watermark in the image

Style: Photorealistic, magazine-quality food photography that would be used in high-end restaurant menu or food advertising.
"""
    try:
        gen_params = {
            "model": "gpt-image-1",
            "prompt": food_prompt,
            "size": "1024x1024",
            "quality": "high",
            "n": 1
        }

        result = client.images.generate(**gen_params)

        if not result or not result.data or len(result.data) == 0:
            st.error(f"No image data returned for {dish_name}")
            return None, None

        # Get base64 image data directly ( it should work directly )
        image_base64 = result.data[0].b64_json

        # Check if base64 data is valid
        if not image_base64:
            st.error(f"No base64 image returned for {dish_name}")
            return None, None

        image_bytes = base64.b64decode(image_base64)

        # Create a data URL for display purpose
        image_url = f"data:image/jpeg;base64,{image_base64}"

        return image_url, image_bytes
    except Exception as e:
        st.error(f"Error generating food image for {dish_name}: {str(e)}")
        return None, None

# Main UI
st.title("GenAI - Menu Visual Generator")
st.markdown("Transforms your text menu into beautiful visual menu with food photos and nutritional Information!")

if client:
    st.markdown("### Upload your Menu Image")
    st.markdown("Upload a photo of your text-only menu and we'll extract all items automatically!")

    uploaded_file = st.file_uploader(
        "Choose menu image",
        type=['jpg', "jpeg", "png"],
        help="Upload a clear image of your menu with readable text"
    )

    if uploaded_file:
        st.image(uploaded_file, caption="Uploaded Menu Image", width=400)

        if st.button("Extract menu Items", type="primary"):
            with st.spinner("Analyzing menu image and extracting items...."):
                menu_data = create_menu_text_from_image(uploaded_file)

                if "error" in menu_data:
                    st.error(f"Error extracting menu data: {menu_data['error']}")
                    if "raw_text" in menu_data:
                        st.text_area("Raw extracted Text", menu_data['raw_text'], height=200)

                else:
                    st.success("Menu items extracted successfully!")
                    st.session_state.menu_data = menu_data

                    # Display extracted menu structure
                    st.markdown("### Extracted Menu Structure")

                    col1, col2 = st.columns([2, 1])

                    with col1:
                        st.markdown(f"**Restaurant** {menu_data.get('restaurant_image', 'Not detected')}")
                        st.markdown(f"**Cuisine Type** {menu_data.get('cuisine_type', 'Not detected')}")

                        for category, items in menu_data.get('categories', {}).items():
                            st.markdown(f"**{category}**")
                            for item in items:
                                st.markdown(f"- {item['name']} - {item.get('price', 'N/A')}")
                                if item.get('description'):
                                    st.markdown(f"  *{item['description']}*")

                    with col2:
                        st.markdown("**Menu Statistics**")
                        total_items = sum(len(items) for items in menu_data.get('categories', {}).values())
                        st.metric("Total Items", total_items)
                        st.metric("Categories", len(menu_data.get('categories', {})))

    if "menu_data" in st.session_state:
        st.markdown("---")
        st.markdown("## Generate Visual Menu!")

        #configuration options
        cols = st.columns(3)
        col1, col2, col3 = cols[0], cols[1], cols[2]

        with col1:
            image_style = st.selectbox(
                "Food Photo Style",
                [
                    "Professional Food Photography",
                    "Rustic Homestyle",
                    "Modern Minimalstic",
                    "Vibrant Colorful"
                ],
                help="Choose the style of food photography"
            )

        with col2:
            layout_style = st.selectbox(
                "Food Description Display",
                [
                    "Show Descriptions",
                    "Hide Descriptions",
                    "Short Descriptions only"
                ],
                help="Choose how to display food descriptions"
            )

        with col3:
            include_nutrition = st.checkbox(
                "Include Nutritional Information",
                value=True,
                help="Add calories and nutritional information to each dish"
            )


    if st.button("Generate Complete Visual Menu", type="primary", use_container_width=True):
        menu_data = st.session_state.menu_data

        with st.spinner("Generating your visual menu... This may take several minutes as we generate food photos for each item."):

            # progress tracking
            progress_bar = st.progress(0, text="Generating food photos for each item")
            status_text = st.empty()

            dish_images = {}
            nutritional_data = {}

            # get all dishes
            all_dishes = []
            for category, items in menu_data.get('categories', {}).items():
                for item in items:
                    all_dishes.append((category, item))

            total_dishes = len(all_dishes)

            # Generate images and nutrition for each dish
            for i, (category, item) in enumerate(all_dishes):
                dish_name = item['name']
                description = item.get('description', "")

                status_text.text(f"Generating food photo for {dish_name} ({i+1}/{total_dishes})")

                # When you generate the image
                image_bytes, image_url = generate_food_image(
                    dish_name,
                    description,
                    menu_data.get('cuisine_type', "Fine Dining"),
                    image_style
                )

                # Store the image data in the dish_images dictionary
                if image_bytes and image_url:
                    dish_images[dish_name] = {
                        'bytes': image_bytes,
                        'url': image_url
                    }

                # Get nutritional information
                if include_nutrition:
                    nutritional_data[dish_name] = get_nutrinutional_info(dish_name, menu_data.get('cuisine_type', "Fine Dining"))

                # update progress
                progress_bar.progress((i+1)/total_dishes)

            progress_bar.progress(1.0)
            status_text.text("Visual menu generation complete!")

            # Display results
            st.markdown("### Your Visual Menu Results")

            # Show individual food photo
            st.markdown("### Generated Food Photos with Nutritional Information")

            for category, items in menu_data.get('categories', {}).items():
                st.markdown(f"**{category}**")
                cols = st.columns(min(len(items), 3))

                for i, item in enumerate(items):
                    dish_name = item['name']
                    col_index = i % 3

                    with cols[col_index]:
                        if dish_name in dish_images:
                            st.image(
                                dish_images[dish_name]['url'],
                                caption=f"{dish_name} - {item.get('price', '')}",
                                use_column_width=True
                            )

                        if include_nutrition and dish_name in nutritional_data:
                            nutrition = nutritional_data[dish_name]
                            st.markdown(f"**Nutrition** {nutrition.get('calories', 'N/A')} cal, {nutrition.get('protein', 'N/A')} protein")

                        if item.get('description'):
                            st.markdown(f"*{item['description']}*")


        # Download Options
        st.markdown("### Download Individual Food Photos")

        # Create download options for individual images
        download_cols = st.columns(3)
        col_count = 0

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        for category, items in menu_data.get('categories', {}).items():
            st.markdown(f"**Download {category} Photos:**")
            for item in items:
                dish_name = item['name']
                if dish_name in dish_images:
                    clean_name = re.sub(r'[^a-zA-Z0-9]', '', dish_name)

                    with download_cols[col_count % 3]:
                        st.download_button(
                            f"Download {dish_name} Photo",
                            dish_images[dish_name]['bytes'],
                            file_name=f"{clean_name}_{timestamp}.png",
                            mime="image/png",
                            help=f"Download {dish_name} food photo",
                            use_container_width=True
                        )
                    col_count += 1

        restaurant_name_clean = re.sub(r'[^a-zA-Z0-9]', '', menu_data.get('restaurant_name', "menu"))

        st.markdown("### Download Menu Report")

        col_download1, col_download2 = st.columns(2)

        with col_download1:
            st.markdown(f"**Bulk Download**")
            st.info(f"Individual photos available above, or use menu report for compelete details")

        with col_download2:
            menu_summary = f"""VISUAL MENU GENERATED REPORT
Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

Restaurant: {menu_data.get('restaurant_name', 'N/A')}
Cuisine Type: {menu_data.get('cuisine_type', 'N/A')}
Image Style: {image_style}
Food Description Display: {layout_style}
Nutritional Information: {'Included' if include_nutrition else 'Excluded'}

Menu Items Generated: {total_dishes}
"""

            for category, items in menu_data.get('categories', {}).items():
                menu_summary += f"\n{category}:\n"
                for item in items:
                    menu_summary += f"- {item['name']} - {item.get('price', 'N/A')}\n"
                    if include_nutrition and item['name'] in nutritional_data:
                        nutrition = nutritional_data[item['name']]
                        menu_summary += f"Nutrition: {nutrition.get('calories', 'N/A')} cal, {nutrition.get('protein', 'N/A')} protein\n"

            menu_summary += f"""
Generation Statistics:
- Total Items: {sum(len(items) for items in menu_data.get('categories', {}).values())}
- Categories: {len(menu_data.get('categories', {}))}
- Food Photos Generated: {len(dish_images)}
- Processing Time: Several minutes

File Included:
- Individual food photography for each menu item
- Nutritional information (if selected)

Usage Recommendations:
- Use individual food photos for online menus and delivery apps
- Share on social media to showcase specific dishes
- Use for promotional materials and advertisements
- Update photos seasonally or when menu changes
"""

            st.download_button(
                "Download Menu Report",
                menu_summary.encode("utf-8"),
                file_name=f"menu_report_{restaurant_name_clean}_{timestamp}.txt",
                mime="text/plain",
                help="Download detailed generation report",
                use_container_width=True
            )
else:
    st.info("Please enter OpenAI key to get started!.")

    st.markdown("""
    ### Trasform you Menu with AI
    
    ### **Upload Any Menu Photo**
    - Take a picture of your menu and we'll extract all items automatically!
    
    ### **Visualize Your Menu**
    - Generate beautiful visual menu with food photos and nutritional Information!
    
    ### **Download Your Menu Report**
    - Download detailed report of your menu with all details!    
    """)