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

# Load API key
load_dotenv()
API_KEY = os.getenv("GROQ_API_KEY")

# Initialize Groq client
client = Groq(api_key=API_KEY)

# Streamlit App Title
st.set_page_config(page_title="Timetable Generator", layout="centered")
st.title("📅 AI-Powered Timetable Generator")

# Make UI mobile-friendly
st.markdown(
    """
    <style>
        body { font-size: 18px; }
        .stTextArea textarea { font-size: 16px !important; }
        .stDownloadButton button { font-size: 16px !important; padding: 10px; }
        .stButton button { font-size: 16px !important; padding: 10px; }
    </style>
    """,
    unsafe_allow_html=True,
)

# Select Study Field
study_field = st.selectbox("Choose Your Study Field:", ["B.E (Engineering)", "Medical", "Arts"])

# Define subjects for each field
subjects = {
    "B.E (Engineering)": ["Mathematics", "Physics", "Thermodynamics", "Mechanics", "Electronics"],
    "Medical": ["Anatomy", "Physiology", "Pharmacology", "Pathology", "Surgery"],
    "Arts": ["History", "Literature", "Philosophy", "Music", "Fine Arts"]
}

# Days of the week
week_days = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday"]

# Generate Timetable Button
if st.button("Generate Timetable"):
    user_prompt = f"Create a structured weekly timetable for {study_field} students with subjects: {', '.join(subjects[study_field])}. Assign subjects from Monday to Friday."

    try:
        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": user_prompt}]
        )
        timetable = response.choices[0].message.content

        # Display Timetable
        st.write("### 📌 AI-Generated Timetable:")
        st.write(timetable)

        # Function to generate PDF
        def generate_pdf(timetable_text):
            pdf = FPDF()
            pdf.set_auto_page_break(auto=True, margin=10)
            pdf.add_page()
            pdf.set_font("Arial", size=12)
            pdf.cell(200, 10, "AI-Generated Weekly Timetable", ln=True, align="C")
            pdf.ln(10)
            for line in timetable_text.split("\n"):
                pdf.multi_cell(0, 10, line)
            return pdf

        # Generate and download PDF
        if st.button("📥 Download Timetable as PDF"):
            pdf = generate_pdf(timetable)
            pdf_filename = "Weekly_Timetable.pdf"
            pdf.output(pdf_filename)
            with open(pdf_filename, "rb") as file:
                st.download_button(label="📥 Download PDF", data=file, file_name=pdf_filename, mime="application/pdf")

    except Exception as e:
        st.error(f"Error: {e}")