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']

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

def initialize_api_client():
    """Initialize Groq API client with environment variables"""
    load_dotenv()
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        st.error("API key not found. Please check the .env configuration.")
        st.stop()
    try:
        return Groq(api_key=api_key)
    except Exception as e:
        st.error(f"Error initializing API client: {e}")
        st.stop()

def encode_image(image_path):
    """Convert image file 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"Error: Image file '{image_path}' not found.")
        return ""
    except Exception as e:
        st.error(f"Error encoding image '{image_path}': {e}")
        return ""

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

def generate_pdf(report_text, logo_b64):
    """Generate PDF report with nutrition analysis and logo"""
    try:
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=letter)
        styles = getSampleStyleSheet()
        
        # Decode logo image from base64 and resize
        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  # Adjust as needed
        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)
        
        # Build PDF content
        story = [
            logo,
            Spacer(1, 12),
            Paragraph("<b>Nutrition Analysis Report</b>", styles['Title']),
            Spacer(1, 12),
            Paragraph(report_text.replace('\n', '<br/>'), styles['BodyText'])
        ]
        
        doc.build(story)
        buffer.seek(0)
        return buffer
    except Exception as e:
        st.error(f"Error generating PDF: {e}")
        return None

def generate_analysis(uploaded_file, client):
    """Generate nutrition analysis using AI (Groq API)"""
    base64_image, img_format = process_image(uploaded_file)
    if not base64_image:
        st.error("Failed to process image. Cannot generate analysis.")
        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=[{
                "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.
                """
            }],
            temperature=0.2,
            max_tokens=400,
            top_p=0.5
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"Error communicating with API: {e}")
        return None

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

def display_main_interface():
    """Render the primary user interface"""
    logo_b64 = encode_image("src/logo.png")
    
    # Display logo and title
    if not logo_b64:
        st.error("Logo image not found or failed to encode.")
        return
    
    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("---")
    
    # Display analysis results if available
    if st.session_state.get('analysis_result'):
        col1, col2 = st.columns([1, 1])
        
        # Column for download button
        with col1:
            pdf_report = generate_pdf(st.session_state.analysis_result, logo_b64)
            if pdf_report:
                st.download_button("📄 Download Nutrition Report", data=pdf_report, file_name="nutrition_report.pdf", mime="application/pdf")
            else:
                st.error("Failed to generate PDF report.")
        
        # Column for clear button
        with col2:
            if st.button("Clear Analysis 🗑️"):
                st.session_state.pop('analysis_result', None)
                st.rerun()
    
    if st.session_state.get('analysis_result'):
        st.markdown("### 🎯 Nutrition Analysis Report")
        st.info(st.session_state.analysis_result)

def render_sidebar(client):
    """Render the sidebar with image upload and analysis functionality"""
    with st.sidebar:
        st.subheader("Image Upload")
        uploaded_file = st.file_uploader("Upload Food Image", type=ALLOWED_FILE_TYPES)
        
        if uploaded_file:
            st.image(Image.open(uploaded_file), caption="Uploaded Food Image")
            if st.button("Analyze Meal 🍽️"):
                with st.spinner("Analyzing image..."):
                    report = generate_analysis(uploaded_file, client)
                    if report:
                        st.session_state.analysis_result = report
                        st.rerun()
                    else:
                        st.error("Failed to generate analysis. Please try again.")

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

def main():
    """Main application controller"""
    client = initialize_api_client()
    display_main_interface()
    render_sidebar(client)

if __name__ == "__main__":
    main()