import json
import requests
import os
import streamlit as st
import pandas as pd
import openpyxl
import torch
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# ✅ 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
st.sidebar.markdown("## 📂 Upload Your Timetable Files")
uploaded_master = st.sidebar.file_uploader("Upload Master Timetable", type=["xlsx"])
uploaded_lab = st.sidebar.file_uploader("Upload Lab Timetable", type=["xlsx"])
uploaded_classroom = st.sidebar.file_uploader("Upload Classroom Timetable", type=["xlsx"])
uploaded_individual = st.sidebar.file_uploader("Upload Individual Timetable", type=["xlsx"])

uploaded_files = {
    "Master Timetable": uploaded_master,
    "Lab Timetable": uploaded_lab,
    "Classroom Timetable": uploaded_classroom,
    "Individual Timetable": uploaded_individual,
}

# ✅ Load Timetable Data (Directly from Uploaded File)
def load_timetable(file):
    if not file:
        return None
    wb = openpyxl.load_workbook(file)
    sheet = wb.active
    return [row for row in sheet.iter_rows(values_only=True)]

# ✅ Ask Mistral AI via API
def ask_mistral_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": "mistralai/Mistral-7B-Instruct-v0.2",
        "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
def auto_schedule(file):
    if not file:
        return "No timetable uploaded."

    wb = openpyxl.load_workbook(file)
    sheet = wb.active

    empty_slots = []
    for row_idx, row in enumerate(sheet.iter_rows(min_row=2, values_only=True), start=2):
        if None in row or "" in row:
            empty_slots.append(row_idx)

    for row_idx in empty_slots:
        query = f"Suggest a subject and faculty for the empty slot in row {row_idx}."
        suggestion = ask_mistral_api(query)

        try:
            subject, faculty = suggestion.split(", Faculty: ")
            sheet.cell(row=row_idx, column=4, value=subject.strip())
            sheet.cell(row=row_idx, column=5, value=faculty.strip())
        except:
            continue  

    return f"Auto-scheduling completed for {len(empty_slots)} slots."

# ✅ AI Query Section
st.markdown("## 🤖 Ask Mistral 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_mistral_api(user_query)
    st.write("🧠 **Mistral AI 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"):
    result = auto_schedule(uploaded_files[selected_file])
    st.write("✅", result)

# ✅ Display Uploaded Timetables
st.markdown("## 📜 View Uploaded Timetables")

for name, file in uploaded_files.items():
    if file:
        df = pd.read_excel(file)
        st.markdown(f"### {name}")
        st.dataframe(df)

# ✅ Inject JavaScript for Real-Time Chat
st.markdown("""
    <script>
        async function fetchChatResponse() {
            const apiKey = document.getElementById("hf-api-key").value;
            const userInput = document.getElementById("user-input").value;
            if (!apiKey) {
                alert("Please enter your Hugging Face API key.");
                return;
            }

            try {
                const response = await fetch("https://api-inference.huggingface.co/v1/chat/completions", {
                    method: "POST",
                    headers: {
                        "Authorization": `Bearer ${apiKey}`,
                        "Content-Type": "application/json"
                    },
                    body: JSON.stringify({
                        model: "mistralai/Mistral-7B-Instruct-v0.2",
                        messages: [{ role: "user", content: userInput }],
                        max_tokens: 500
                    })
                });

                if (!response.ok) {
                    throw new Error(`API Error: ${response.statusText}`);
                }

                const data = await response.json();
                const botMessage = data.choices[0].message.content;

                document.getElementById("chat-box").innerHTML += 
                    `<div class='bot-message'><strong>VarunGPT-3:</strong> ${botMessage}</div>`;

            } catch (error) {
                console.error("Error fetching chat response:", error);
                document.getElementById("chat-box").innerHTML += 
                    `<div class='bot-message'><strong>Error:</strong> Unable to fetch response.</div>`;
            }
        }
    </script>
""", unsafe_allow_html=True)

# ✅ Chat UI with User Input for API Key
st.markdown("""
    <div style="text-align: center;">
        <input id="hf-api-key" type="password" placeholder="Enter Hugging Face API Key" 
               style="width: 50%; padding: 8px; margin-bottom: 10px;"/>
        <br/>
        <input id="user-input" type="text" placeholder="Type your message..." 
               style="width: 50%; padding: 8px;"/>
        <button onclick="fetchChatResponse()" 
                style="padding: 10px 20px; background-color: #4CAF50; color: white; border: none; cursor: pointer;">
            Ask AI
        </button>
        <div id="chat-box" style="margin-top: 20px; text-align: left;"></div>
    </div>
""", unsafe_allow_html=True)