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 resize_image(image):
    """Resize the image to reduce size before encoding it to base64"""
    max_width = 800  # Max width in pixels
    max_height = 800  # Max height in pixels
    image.thumbnail((max_width, max_height))
    return image


def compress_image(image):
    """Compress the image to JPEG format with a lower quality to reduce size"""
    buffer = io.BytesIO()
    image.save(buffer, format="JPEG", quality=70)  # Reduce quality to 70%
    return base64.b64encode(buffer.getvalue()).decode('utf-8')


def process_image(uploaded_file):
    """Convert uploaded image file to base64 string with compression"""
    try:
        image = Image.open(uploaded_file)
        image = resize_image(image)  # Resize the image to reduce its size
        base64_image = compress_image(image)  # Compress the image
        return base64_image, "jpeg"  # Use 'jpeg' as format
    except Exception as e:
        st.error(f"Error processing image: {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 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=[{
                "role": "system",  # Define the role for the system message
                "content": """
                    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.
                """
            },
            {
                "role": "user",  # Define the role for the user message
                "content": f"Analyze this image and provide the nutrition analysis: {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"Error communicating with API: {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()