import json
import requests
import streamlit as st
import pandas as pd
import openpyxl
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from transformers import pipeline
from docx import Document
from io import BytesIO

# ✅ Streamlit UI Setup
st.set_page_config(page_title="AI-Powered Timetable", layout="wide")
st.markdown("<h1 style='text-align: center; color: #4CAF50;'>📅 AI-Powered Timetable</h1>", unsafe_allow_html=True)

# ✅ API Key Input
st.sidebar.markdown("## 🔑 Enter Hugging Face API Key")
hf_api_key = st.sidebar.text_input("API Key", type="password")

# ✅ File Upload Section (Excel & DOCX)
st.sidebar.markdown("## 📂 Upload Your Timetable Files")
uploaded_files = {
    "Master Timetable": st.sidebar.file_uploader("Upload Master Timetable", type=["xlsx", "docx"]),
    "Lab Timetable": st.sidebar.file_uploader("Upload Lab Timetable", type=["xlsx", "docx"]),
    "Classroom Timetable": st.sidebar.file_uploader("Upload Classroom Timetable", type=["xlsx", "docx"]),
    "Individual Timetable": st.sidebar.file_uploader("Upload Individual Timetable", type=["xlsx", "docx"]),
}

# ✅ Load DOCX File Function
def load_docx(file):
    doc = Document(file)
    return "\n".join([para.text for para in doc.paragraphs])

# ✅ Save DOCX File Function
def save_docx(text):
    doc = Document()
    for line in text.split("\n"):
        doc.add_paragraph(line)
    output = BytesIO()
    doc.save(output)
    output.seek(0)
    return output

# ✅ Load XLSX File
def load_xlsx(file):
    return pd.read_excel(file)

# ✅ Ask TinyLlama AI via API
def ask_tinyllama_api(query):
    if not hf_api_key:
        return "Error: Please enter your API key."

    url = "https://api-inference.huggingface.co/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {hf_api_key}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        "messages": [{"role": "user", "content": query}],
        "max_tokens": 500
    }

    response = requests.post(url, headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return f"API Error: {response.status_code} - {response.text}"

# ✅ Auto-Schedule Missing Slots in Excel
def auto_schedule(file):
    if not file:
        return "No timetable uploaded."

    df = pd.read_excel(file)
    empty_slots = df[df.isnull().any(axis=1)]

    for index, row in empty_slots.iterrows():
        query = f"Suggest a subject and faculty for an empty slot on {row.get('Day', 'Unknown')} at {row.get('Time', 'Unknown')}."
        suggestion = ask_tinyllama_api(query)
        try:
            subject, faculty = suggestion.split(", Faculty: ")
            df.loc[index, "Subject"] = subject.strip()
            df.loc[index, "Faculty"] = faculty.strip()
        except:
            continue  

    output = BytesIO()
    df.to_excel(output, index=False, engine='openpyxl')
    output.seek(0)
    return output

# ✅ AI Query Section
st.markdown("## 🤖 Ask TinyLlama AI About Your Timetable")
user_query = st.text_input("Type your question here (e.g., 'Who is free at 10 AM on Monday?')")

if st.button("Ask AI via API"):
    ai_response = ask_tinyllama_api(user_query)
    st.write("🧠 **TinyLlama Suggests:**", ai_response)

# ✅ Auto-Schedule Feature
st.markdown("## 📅 Auto-Schedule Missing Timetable Slots")
selected_file = st.selectbox("Choose a timetable file to auto-fill missing slots:", list(uploaded_files.keys()))

if st.button("Auto-Schedule"):
    if uploaded_files[selected_file]:
        scheduled_file = auto_schedule(uploaded_files[selected_file])
        st.download_button("📥 Download Auto-Scheduled Timetable", scheduled_file, file_name=f"{selected_file}_auto.xlsx", mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
    else:
        st.write("⚠️ Please upload a file first.")

# ✅ Display & Edit Uploaded Timetables (DOCX + XLSX)
st.markdown("## 📜 View & Edit Uploaded Timetables")

for name, file in uploaded_files.items():
    if file:
        if file.name.endswith(".xlsx"):
            df = load_xlsx(file)
            edited_df = st.data_editor(df, num_rows="dynamic")
            csv_output = edited_df.to_csv(index=False).encode()
            st.download_button("📥 Download Edited Excel", csv_output, file_name=f"{name}.csv", mime="text/csv")

        elif file.name.endswith(".docx"):
            doc_text = load_docx(file)
            edited_text = st.text_area(f"📄 Edit {name}:", doc_text, height=400)
            if st.button(f"💾 Save & Download {name}"):
                docx_output = save_docx(edited_text)
                st.download_button("📥 Download Edited DOCX", docx_output, file_name=f"{name}.docx", mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document")

# ✅ Export to PDF
st.markdown("## 📄 Export to PDF")
if st.button("Export to PDF"):
    pdf_output = BytesIO()
    c = canvas.Canvas(pdf_output, pagesize=letter)
    c.drawString(100, 750, "AI-Powered Timetable Export")
    y_position = 730
    for name, file in uploaded_files.items():
        if file:
            c.drawString(100, y_position, f"{name}: {file.name}")
            y_position -= 20
    c.save()
    pdf_output.seek(0)
    st.download_button("📥 Download PDF", pdf_output, file_name="timetable.pdf", mime="application/pdf")