cimport os
import streamlit as st
from groq import Groq
from fpdf import FPDF

# Set up Groq API client
groq_api_key = os.getenv("GROQ_API_KEY")
client = Groq(api_key=groq_api_key)

# Initialize Streamlit app
st.title("Timetable Generator")

# Step 1: Collect user inputs for timetable generation
subjects = ['Mathematics', 'Chemistry', 'Science', 'Engineering']
days_of_week = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']

num_days = st.slider('Select the number of days for your timetable', 1, 7, 5)
time_slots = st.slider('Select the number of time slots per day', 1, 6, 4)

# User input for subjects and time slots
subject_schedule = {day: {subject: [] for subject in subjects} for day in days_of_week[:num_days]}

for day in days_of_week[:num_days]:
    st.subheader(f"Schedule for {day}")
    for subject in subjects:
        subject_schedule[day][subject] = st.multiselect(f"Choose time slots for {subject} on {day}",
                                                       options=[f"{hour}:00" for hour in range(8, 8 + time_slots)],
                                                       label_visibility="collapsed")

# Step 2: Request timetable generation from Groq API
if st.button("Generate Timetable"):
    timetable_input = f"Create a timetable for {subject_schedule}. The timetable should span {num_days} days and include the following subjects: {', '.join(subjects)}."

    # API call to Groq's Llama model
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": timetable_input}],
        model="llama-3.3-70b-versatile"
    )
    
    # Extract generated timetable
    generated_timetable = chat_completion.choices[0].message.content
    st.write(generated_timetable)

    # Step 3: Generate a PDF with the timetable
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()

    # Set Title
    pdf.set_font("Arial", size=12, style='B')
    pdf.cell(200, 10, txt="Weekly Timetable", ln=True, align='C')

    # Add each day's schedule to the PDF
    pdf.set_font("Arial", size=10)
    for day in subject_schedule:
        pdf.ln(10)
        pdf.cell(200, 10, txt=day, ln=True, align='L')
        for subject in subject_schedule[day]:
            time_slots = ", ".join(subject_schedule[day][subject])
            pdf.cell(200, 10, txt=f"{subject}: {time_slots}", ln=True)

    # Save the PDF
    pdf.output("timetable.pdf")

    # Step 4: Allow user to download the PDF
    with open("timetable.pdf", "rb") as f:
        st.download_button(
            label="Download Timetable as PDF",
            data=f,
            file_name="timetable.pdf",
            mime="application/pdf"
        )