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

# Set the Groq API key from Hugging Face Secrets (ensure to add it in your Space's settings)
api_key = os.getenv("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
")
client = Groq(api_key=api_key)

# Function to generate a timetable
def generate_timetable(teachers, subjects, slots):
    """
    Generate a timetable using the Groq API.
    """
    try:
        prompt = f"""
        Create a weekly timetable given the following data:
        Teachers: {", ".join(teachers)}
        Subjects: {", ".join(subjects)}
        Available Slots: {", ".join(slots)}.
        """
        
        chat_completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama3-8b-8192",
        )
        
        return chat_completion.choices[0].message.content

    except Exception as e:
        st.error(f"Error generating timetable: {e}")
        return None

# Function to save timetable as PDF
def save_timetable_as_pdf(timetable, filename="timetable.pdf"):
    """
    Save the generated timetable as a PDF.
    """
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    
    pdf.cell(200, 10, txt="Class Timetable", ln=True, align="C")
    pdf.ln(10)  # Add some space
    
    for line in timetable.split("\n"):
        pdf.cell(200, 10, txt=line, ln=True)
    
    pdf.output(filename)
    return filename

# Streamlit App UI
st.title("Class Timetable Generator")
st.write("Generate and manage class schedules effortlessly.")

# Input fields
teachers = st.text_input("Enter the names of teachers (comma-separated)").split(",")
subjects = st.text_input("Enter the subjects (comma-separated)").split(",")
slots = st.text_input("Enter the available slots (comma-separated)").split(",")

if st.button("Generate Timetable"):
    if teachers and subjects and slots:
        st.info("Generating timetable...")
        timetable = generate_timetable(teachers, subjects, slots)
        
        if timetable:
            st.success("Timetable generated successfully!")
            st.text_area("Generated Timetable", timetable, height=300)

            # Option to download the timetable as a PDF
            if st.button("Download PDF"):
                pdf_file = save_timetable_as_pdf(timetable)
                with open(pdf_file, "rb") as file:
                    st.download_button("Download Timetable PDF", file, file_name="timetable.pdf")
        else:
            st.error("Failed to generate the timetable. Please check the inputs or API settings.")
    else:
        st.warning("Please fill all fields to generate the timetable.")