import streamlit as st
import os
import io
import tempfile
from PIL import Image
from fpdf import FPDF

# --- Page Setup ---
st.set_page_config(page_title="Al Dhafra Photo Generator", layout="wide")

# Constants
LOGO_PATH = "Al Dhafra Logo.png"
GRADE_LIST = ["KG1", "KG2"] + [f"Grade {i}" for i in range(1, 13)]
SECTION_LIST = ["A", "B", "C", "D"]
current_year = 2025
AY_LIST = [f"AY {y}-{str(y+1)[2:]}" for y in range(current_year, current_year + 6)]
GOLD_COLOR = (184, 134, 11) 

# --- Loading Logic ---
if 'ai_loaded' not in st.session_state:
    with st.spinner("💫 Loading AI Engine... | جاري تحميل المحرك..."):
        from rembg import remove
        st.session_state['ai_loaded'] = True
        st.session_state['remove_fn'] = remove
else:
    from rembg import remove

# --- UI Tabs ---
tab1, tab2 = st.tabs(["🚀 Generator | المولد", "📖 How to Use | طريقة الاستخدام"])

with tab1:
    st.title("⭐ Al Dhafra Star of the Week")
    
    st.sidebar.header("Photo Settings | إعدادات الصورة")
    selected_ay = st.sidebar.selectbox("Academic Year | السنة الدراسية", AY_LIST)
    selected_term = st.sidebar.selectbox("Select Term | اختر الفصل الدراسي", ["Term 1", "Term 2", "Term 3"])
    selected_week = st.sidebar.selectbox("Select Week | اختر الأسبوع", [f"Week {i}" for i in range(1, 13)])
    border_color = st.sidebar.color_picker("Border Color | لون الإطار", "#001F3F")

    h_color = border_color.lstrip('#')
    rgb_border = tuple(int(h_color[i:i+2], 16) for i in (0, 2, 4))

    num_students = st.number_input("Number of students | عدد الطلاب:", min_value=1, step=1)
    students_data = []

    for i in range(num_students):
        with st.expander(f"Student {i+1} | الطالب {i+1}", expanded=True):
            col1, col2 = st.columns([1, 1])
            with col1:
                name = st.text_input(f"Full Name | الاسم الكامل", key=f"n{i}")
                sub_col1, sub_col2 = st.columns(2)
                with sub_col1:
                    grade = st.selectbox(f"Grade | الصف", GRADE_LIST, key=f"g{i}")
                with sub_col2:
                    section = st.selectbox(f"Section | الشعبة", SECTION_LIST, key=f"s{i}")
            with col2:
                img = st.file_uploader(f"Upload Photo | رفع الصورة", type=['jpg', 'png', 'jpeg'], key=f"i{i}")
                
            if name and img:
                students_data.append({"name": name, "grade": grade, "section": section, "image": img})

    def create_pdf(data, border_rgb, term, week, ay):
        pdf = FPDF()
        for student in data:
            pdf.add_page()
            
            # 1. Border
            pdf.set_draw_color(*border_rgb)
            pdf.set_line_width(5)
            pdf.rect(10, 10, 190, 277)
            
            if os.path.exists(LOGO_PATH):
                pdf.image(LOGO_PATH, x=170, y=12, w=20) 
            
            # 2. Header (Term | Week | AY)
            pdf.set_xy(15, 15)
            pdf.set_font("Arial", 'B', 11)
            pdf.set_text_color(160, 160, 160)
            pdf.cell(0, 10, f"{term} | {week} | {ay}", ln=True)

            # 3. AI Processing (Background Removal & Smart Crop)
            raw_img = Image.open(student["image"])
            processed_img = st.session_state['remove_fn'](raw_img, alpha_matting=True)
            bbox = processed_img.getbbox()
            if bbox:
                processed_img = processed_img.crop(bbox)

            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
                processed_img.save(tmp.name, format='PNG')
                temp_path = tmp.name

            # 4. Golden Title
            pdf.set_y(28) 
            pdf.set_font("Arial", 'B', 32)
            pdf.set_text_color(*GOLD_COLOR) 
            pdf.cell(200, 20, "STAR OF THE WEEK", ln=True, align='C')

            # 5. Combined Single Line: NAME - GRADE - SECTION
            # We use a slightly smaller font (28) to ensure long names fit on one line
            pdf.set_y(48)
            pdf.set_font("Arial", 'B', 28)
            pdf.set_text_color(*border_rgb)
            combined_text = f"{student['name'].upper()} - {student['grade']}-{student['section']}"
            pdf.cell(200, 15, combined_text, ln=True, align='C')

            # 6. Student Photo (Proportional Scaling & Maximum Space)
            # Photo now starts at y=65 since we saved space by combining lines
            pdf.image(temp_path, x=35, y=65, w=140, h=0)
            os.unlink(temp_path)
            
        return pdf.output(dest='S').encode('latin-1')

    if st.button("✨ Generate Star of the Week Photos | إنشاء الصور"):
        if not students_data:
            st.warning("Please fill details | يرجى ملء البيانات")
        else:
            with st.spinner("🎨 Creating... | جاري الإنشاء..."):
                pdf_out = create_pdf(students_data, rgb_border, selected_term, selected_week, selected_ay)
                
            st.balloons()
            st.success("✅ Ready! | جاهز")
            st.download_button(
                label="📥 Download Photos PDF | تحميل الملف",
                data=pdf_out,
                file_name=f"Star_of_the_Week_{selected_ay}.pdf",
                mime="application/pdf",
                use_container_width=True
            )