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"""
    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:
        return ""


def process_image(uploaded_file):
    """Convert image 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"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()
    
    # Decode the base64 logo image
    logo_data = base64.b64decode(logo_b64)
    logo_image = Image.open(io.BytesIO(logo_data))
    
    # Resize the logo to fit the page width (you can adjust size if necessary)
    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)
    
    # Create a ReportLab Image element to add the logo to the PDF
    logo = ReportLabImage(io.BytesIO(logo_data), width=logo_width, height=logo_height)
    
    # Build the PDF content
    story = [
        logo,  # Add the logo at the top of the page
        Spacer(1, 12),  # Space after the logo
        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


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=[
                {
                    "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"API communication error: {e}")
        return None

# UI COMPONENTS
def display_main_interface():
    """Render primary application interface"""
    logo_b64 = encode_image("src/logo.png")
    
    # HTML with inline styles to change text colors
    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'):
        # Create two columns: one for download and one for clear button
        col1, col2 = st.columns([1, 1])
        
        # Left column for the Download button
        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")
        
        # Right column for the Clear button
        with col2:
            if st.button("Clear Analysis 🗑️"):
                st.session_state.pop('analysis_result')
                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):
    """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:
            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)
                    st.session_state.analysis_result = report
                    st.rerun()

# APPLICATION ENTRYPOINT
def main():
    """Primary application controller"""
    client = initialize_api_client()
    display_main_interface()
    render_sidebar(client)

if __name__ == "__main__":
    main()