import streamlit as st
from PIL import Image
import cv2
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
from reportlab.lib.styles import getSampleStyleSheet

# ======================
# CONFIGURATION SETTINGS
# ======================
PAGE_CONFIG = {
    "page_title": "Rice Quality Analyzer",
    "page_icon": "🌾",
    "layout": "wide",
    "initial_sidebar_state": "expanded"
}

ALLOWED_FILE_TYPES = ['png', 'jpg', 'jpeg']
ALLOWED_VIDEO_TYPES = ['mp4', 'avi', 'mov']

CSS_STYLES = """
<style>
    .main { background-color: #f4f9f9; color: #000000; }
    .sidebar .sidebar-content { background-color: #d1e7dd; }
    .stTextInput textarea { color: #000000 !important; }
    .stButton>button { 
        background-color: #21eeef; 
        color: white; 
        font-size: 16px; 
        border-radius: 5px; 
    }
    .report-container {
        background-color: #ffffff;
        border-radius: 15px;
        padding: 20px;
        margin-top: 20px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        border-left: 5px solid #21eeef;
    }
    .report-text {
        font-family: 'Arial', sans-serif;
        font-size: 14px;
        line-height: 1.6;
        color: #2c3e50;
    }
</style>
"""

# ======================
# CORE FUNCTIONS
# ======================
def configure_application():
    """Initialize application settings and styling"""
    st.set_page_config(**PAGE_CONFIG)
    st.markdown(CSS_STYLES, unsafe_allow_html=True)

def initialize_api_client():
    """Create and validate 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 process_image_data(uploaded_file):
    """Convert image to base64 encoded 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: {str(e)}")
        return None, None

def extract_video_frames(uploaded_video):
    """Extract frames from uploaded video for analysis"""
    try:
        tfile = io.BytesIO(uploaded_video.read())
        temp_filename = "temp_video.mp4"
        with open(temp_filename, "wb") as f:
            f.write(tfile.getvalue())
        
        cap = cv2.VideoCapture(temp_filename)
        frame_list = []
        frame_count = 0
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret or frame_count > 10:  # Process only up to 10 frames
                break
            frame_list.append(frame)
            frame_count += 1
        
        cap.release()
        os.remove(temp_filename)
        return frame_list
    except Exception as e:
        st.error(f"Video processing error: {str(e)}")
        return []

def generate_pdf_report(report_text):
    """Generate PDF document from report text"""
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter)
    styles = getSampleStyleSheet()
    story = []
    
    title = Paragraph("<b>Rice Quality Report</b>", styles['Title'])
    story.append(title)
    story.append(Spacer(1, 12))
    
    content = Paragraph(report_text.replace('\n', '<br/>'), styles['BodyText'])
    story.append(content)
    
    doc.build(story)
    buffer.seek(0)
    return buffer

def generate_rice_report(image_data, img_format, client):
    """Generate AI-powered rice quality analysis"""
    if not image_data:
        return None
    
    image_url = f"data:image/{img_format.lower()};base64,{image_data}"
    
    try:
        response = client.chat.completions.create(
            model="llama-3.2-11b-vision-preview",
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": (
                        "Analyze the rice grain image and provide a detailed report including:\n"
                        "1. Rice type classification\n"
                        "2. Quality assessment (broken grains %, discoloration %, impurities %)\n"
                        "3. Foreign object detection\n"
                        "4. Size and shape consistency\n"
                        "5. Recommendations for processing or improvement"
                    )},
                    {"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: {str(e)}")
        return None

# ======================
# UI COMPONENTS
# ======================
def display_main_interface():
    """Render primary application interface"""
    st.title("🌾 Rice Quality Analyzer")
    st.subheader("AI-Powered Rice Grain Inspection")
    st.markdown("---")

def render_sidebar(client):
    """Create sidebar interface elements"""
    with st.sidebar:
        st.subheader("Upload Image or Video")
        uploaded_file = st.file_uploader("Select an image or video", type=ALLOWED_FILE_TYPES + ALLOWED_VIDEO_TYPES)
        
        if uploaded_file:
            file_type = uploaded_file.type
            if "video" in file_type:
                frames = extract_video_frames(uploaded_file)
                if frames:
                    st.image(frames[0], caption="Extracted Frame for Analysis", use_column_width=True)
                    frame_img = Image.fromarray(cv2.cvtColor(frames[0], cv2.COLOR_BGR2RGB))
                    buffer = io.BytesIO()
                    frame_img.save(buffer, format="JPEG")
                    base64_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
                    report = generate_rice_report(base64_image, "jpeg", client)
            else:
                st.image(Image.open(uploaded_file), caption="Uploaded Image", use_column_width=True)
                base64_image, img_format = process_image_data(uploaded_file)
                report = generate_rice_report(base64_image, img_format, client)
            
            if report:
                st.markdown("### 📋 Analysis Report")
                st.markdown(report)

# ======================
# APPLICATION ENTRYPOINT
# ======================
def main():
    configure_application()
    groq_client = initialize_api_client()
    display_main_interface()
    render_sidebar(groq_client)

if __name__ == "__main__":
    main()