import streamlit as st
from PIL import Image
import os
import base64
import io
from dotenv import load_dotenv
from groq import Groq
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image as ReportLabImage
from reportlab.lib.styles import getSampleStyleSheet

# ======================
# CONFIGURATION SETTINGS
# ======================
st.set_page_config(
    page_title="Smart Diet Analyzer",
    page_icon="🍎",
    layout="wide",
    initial_sidebar_state="expanded"
)

ALLOWED_FILE_TYPES = ['png', 'jpg', 'jpeg']

# ======================
# RERUN HELPER FUNCTION
# ======================
def rerun():
    """
    Helper function to rerun the app.
    Tries to use st.experimental_rerun if available; otherwise, does nothing.
    """
    if hasattr(st, "experimental_rerun"):
        st.experimental_rerun()
    else:
        # Fallback: you might consider raising an exception to force a rerun.
        # However, this is not recommended for production. Instead, you might simply
        # notify the user to manually refresh the page.
        st.warning("Please refresh the page to update the results.")

# ======================
# UTILITY FUNCTIONS
# ======================

def initialize_api_client():
    """Initialize Groq API client."""
    load_dotenv()
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        st.error("API key not found. Please verify .env configuration.")
        st.stop()
    return Groq(api_key=api_key)


def encode_image(image_path):
    """Encode an image to base64."""
    try:
        with open(image_path, "rb") as img_file:
            return base64.b64encode(img_file.read()).decode("utf-8")
    except FileNotFoundError:
        st.error(f"Logo file not found at {image_path}")
        return ""


def process_image(uploaded_file):
    """Convert image to base64 string."""
    try:
        image = Image.open(uploaded_file)
        buffer = io.BytesIO()
        fmt = image.format if image.format else "PNG"
        image.save(buffer, format=fmt)
        return base64.b64encode(buffer.getvalue()).decode('utf-8'), fmt
    except Exception as e:
        st.error(f"Image processing error: {e}")
        return None, None


def generate_pdf(report_text, logo_b64):
    """Generate a PDF report with logo."""
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter)
    styles = getSampleStyleSheet()
    
    story = []
    
    # If a logo is provided, decode and add it
    if logo_b64:
        try:
            logo_data = base64.b64decode(logo_b64)
            logo_image = Image.open(io.BytesIO(logo_data))
            logo_width, logo_height = logo_image.size
            logo_aspect = logo_height / logo_width
            max_logo_width = 150
            logo_width = min(logo_width, max_logo_width)
            logo_height = int(logo_width * logo_aspect)
            
            logo = ReportLabImage(io.BytesIO(logo_data), width=logo_width, height=logo_height)
            story.append(logo)
            story.append(Spacer(1, 12))
        except Exception as e:
            st.error(f"Error adding logo to PDF: {e}")
    
    story.extend([
        Paragraph("<b>Nutrition Analysis Report</b>", styles['Title']),
        Spacer(1, 12),
        Paragraph(report_text.replace('\n', '<br/>'), styles['BodyText'])
    ])
    
    try:
        doc.build(story)
    except Exception as e:
        st.error(f"Error generating PDF: {e}")
    
    buffer.seek(0)
    return buffer


def generate_analysis(uploaded_file, client):
    """Generate AI-powered food analysis."""
    base64_image, img_format = process_image(uploaded_file)
    if not base64_image:
        return None
    
    image_url = f"data:image/{img_format.lower()};base64,{base64_image}"
    
    try:
        response = client.chat.completions.create(
            model="llama-3.2-11b-vision-preview",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": """
You are an expert nutritionist with advanced image analysis capabilities.  
Your task is to analyze the provided image, identify all visible food items, and estimate their calorie content with high accuracy.  
**Instructions:**  
- Identify and list each food item visible in the image.  
- For each item, estimate the calorie content based on standard nutritional data, considering portion size, cooking method, and food density.  
- Clearly mark any calorie estimate as "approximate" if based on assumptions due to unclear details.  
- Calculate and provide the total estimated calories for the entire meal.  
**Output Format:**  
- Food Item 1: [Name] – Estimated Calories: [value] kcal  
- Food Item 2: [Name] – Estimated Calories: [value] kcal  
- ...  
- **Total Estimated Calories:** [value] kcal  

If the image lacks sufficient detail or is unclear, specify the limitations and include your confidence level in the estimate as a percentage.
                        """},
                        {"type": "image_url", "image_url": {"url": image_url}}
                    ]
                }
            ],
            temperature=0.2,
            max_tokens=400,
            top_p=0.5
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"API communication error: {e}")
        return None

# ======================
# UI COMPONENTS
# ======================

def display_main_interface(logo_b64):
    """Render primary application interface."""
    st.markdown(f"""
        <div style="text-align: center;">
            <img src="data:image/png;base64,{logo_b64}" width="100">
            <h2 style="color: #4CAF50;">Smart Diet Analyzer</h2>
            <p style="color: #FF6347;">AI-Powered Food & Nutrition Analysis</p>
        </div>
    """, unsafe_allow_html=True)
    
    st.markdown("---")
    
    if st.session_state.get('analysis_result'):
        col1, col2 = st.columns(2)
        
        with col1:
            pdf_report = generate_pdf(st.session_state.analysis_result, logo_b64)
            st.download_button("📄 Download Nutrition Report",
                               data=pdf_report,
                               file_name="nutrition_report.pdf",
                               mime="application/pdf")
        
        with col2:
            if st.button("Clear Analysis 🗑️"):
                st.session_state.pop('analysis_result')
                rerun()
    
    if st.session_state.get('analysis_result'):
        st.markdown("### 🎯 Nutrition Analysis Report")
        st.info(st.session_state.analysis_result)


def render_sidebar(client, logo_b64):
    """Create sidebar UI elements."""
    with st.sidebar:
        st.subheader("Image Upload")
        uploaded_file = st.file_uploader("Upload Food Image", type=ALLOWED_FILE_TYPES)
        
        if uploaded_file:
            try:
                image = Image.open(uploaded_file)
                st.image(image, caption="Uploaded Food Image")
            except Exception as e:
                st.error(f"Error displaying image: {e}")
            
            if st.button("Analyze Meal 🍽️"):
                with st.spinner("Analyzing image..."):
                    report = generate_analysis(uploaded_file, client)
                    if report:
                        st.session_state.analysis_result = report
                        rerun()
                    else:
                        st.error("Failed to generate analysis.")


# ======================
# APPLICATION ENTRYPOINT
# ======================

def main():
    """Primary application controller."""
    client = initialize_api_client()
    logo_b64 = encode_image("src/logo.png")
    display_main_interface(logo_b64)
    render_sidebar(client, logo_b64)


if __name__ == "__main__":
    main()