import gradio as gr
import base64
import os
import torch
from transformers import GPT2Tokenizer, AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.llms import HuggingFacePipeline
from langchain.chains import RetrievalQA
from langchain.document_loaders import TextLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# === Load and Embed Documents ===
loader = DirectoryLoader(
    "courses",
    glob="**/*.txt",
    loader_cls=TextLoader
)
raw_docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=100,
    separators=["\n###", "\n##", "\n\n", "\n", ".", " "]
)
docs = text_splitter.split_documents(raw_docs)

embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = Chroma.from_documents(docs, embedding=embedding_model)
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 4})


# === Prompt Template ===
custom_prompt_template = """
You are a helpful and knowledgeable course advisor at Glasgow Caledonian University. Answer the student's question using only the information provided in the context below.
If the context does not contain the answer, politely respond that the information is not available.
Context:
{context}
Question:
{question}
Answer:
"""

prompt = PromptTemplate(
    input_variables=["context", "question"],
    template=custom_prompt_template
)

# === Load Falcon Model ===
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").to("cuda" if torch.cuda.is_available() else "cpu")


generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=256,
    do_sample=False,
    temperature=0.1,
    top_p=0.9
)

llm = HuggingFacePipeline(pipeline=generator, model_kwargs={"return_full_text": False})

# === Setup Retrieval QA Chain ===
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    chain_type="stuff",
    chain_type_kwargs={"prompt": prompt}
)

# === Avatar and Crest ===
avatar_img = "images/logo.jpg"        # Avatar shown beside bot messages
logo = "images/crest.jpg"        # Crest image

# # === Chat Logic with Course Memory ===
def chat_with_bot(message, history, course_state):
    def timestamp():
        return datetime.now().strftime("%H:%M")

    lower_msg = message.lower()
    
    # Try to detect course from first question
    if "msc" in lower_msg:
        course_state = message.strip()  # Store it for later use
        full_query = f"For the course '{course_state}': {message}"
    elif "change course to" in lower_msg:
        course_state = message.replace("change course to", "").strip()
        response = f"🔁 Course changed. Now answering based on: **{course_state}**"
        history.append((f"{message}  \n🕒 {timestamp()}", f"{response}  \n🕒 {timestamp()}"))
        return "", history, course_state
    elif course_state:
        full_query = f"For the course '{course_state}': {message}"
    else:
        full_query = message  # No course memory yet

    try:
        raw_output = qa_chain.run(full_query)
        response = raw_output.split("Answer:")[-1].strip()
        response = response.replace("<|assistant|>", "").strip()
    except Exception as e:
        response = f"⚠️ An error occurred: {str(e)}"

    history.append((f"{message}  \n🕒 {timestamp()}", f"{response}  \n🕒 {timestamp()}"))
    return "", history, course_state




# === Build Gradio UI ===
initial_message = (
    "👋 Welcome! I'm your Assistant for Glasgow Caledoian University.\n"
    "Struggling to find something on our website?\n"
    "Want to know anything about your MSc course?\n\n"
    "Simply ask and we can get started!\n\n"
    "⚠️ Please avoid sharing personal details in this chat.\n"
    "If personal details are ever needed, we’ll always ask for consent first."
)


with gr.Blocks(
    title="🎓 GCU Academic Advisor",
    css="""
    body {
    font-family: 'Segoe UI', sans-serif;
    background-color: #f5f7fa;
    margin: 0;
    padding: 0;
    transition: background-color 0.3s;
}
.dark body {
    background-color: #1e1e2f;
}
.message.user, .message.bot {
    border-radius: 10px;
    padding: 12px 16px;
    animation: fadeIn 0.4s ease-in-out;
}
.message.user {
    background-color: #d2e5ff !important;
    color: #000;
}
.message.bot {
    background-color: #fff !important;
    box-shadow: 0 2px 6px rgba(0,0,0,0.05);
    color: #000;
}
.dark .message.user {
    background-color: #2c3e50 !important;
    color: #ecf0f1;
}
.dark .message.bot {
    background-color: #3c3f58 !important;
    color: #ecf0f1;
}
.gr-button {
    background-color: #0056b3 !important;
    color: #fff !important;
    border-radius: 8px !important;
    padding: 10px 20px !important;
    font-weight: 600;
    transition: background-color 0.2s;
}
.gr-button:hover {
    background-color: #003d80 !important;
}
textarea, input {
    border-radius: 10px !important;
    border: 1px solid #ccc !important;
    padding: 10px !important;
    font-size: 16px !important;
}
.dark textarea, .dark input {
    background-color: #2c2c3c !important;
    color: #fff;
}
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(5px); }
    to { opacity: 1; transform: translateY(0); }
}
.banner-title {
    font-size: 20px;
    margin: 0;
    color: #1a1a1a;
}

.dark .banner-title {
    color: #ffffff;
}

.banner img {
    background-color: white;
}

.dark .banner img {
    background-color: #1e1e2f;
}
    """
) as  demo:

    gr.HTML(""" 
    <script>
       function toggleDarkMode(checked) {
        document.body.classList.toggle('dark', checked);
      }

      window.addEventListener('DOMContentLoaded', () => {
        const observer = new MutationObserver(() => {
          const label = Array.from(document.querySelectorAll('label')).find(el => el.textContent.includes('Dark Mode'));
          if (label && label.querySelector('input[type="checkbox"]')) {
            const checkbox = label.querySelector('input[type="checkbox"]');
            checkbox.addEventListener('change', () => toggleDarkMode(checkbox.checked));
            observer.disconnect();
      }
    });
    observer.observe(document.body, { childList: true, subtree: true });
  });
</script>
""")

    
    with open(logo, "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode("utf-8")

    gr.Markdown(f"""
    <div style="display: flex; align-items: center; gap: 12px; margin-bottom: 10px;" class="banner">
        <img src="data:image/png;base64,{encoded_string}" alt="University Crest" style="height: 40px; border-radius: 6px;">
        <h1 class="banner-title">Glasgow Caledonian University Course Advisor Chatbot</h1>
    </div>
""")

    chatbot = gr.Chatbot(
        avatar_images=(None, avatar_img),
        value=[(initial_message, "I'm ready to help!")],
        show_copy_button=True
    )

    state = gr.State("")


    with gr.Row():
        msg = gr.Textbox(placeholder="Ask a question...", lines=1, scale=5)
        send_btn = gr.Button(" Send", scale=1)
        dark_mode_toggle = gr.Checkbox(label="🌙 Dark Mode", value=False, scale=1)

    def toggle_dark_mode(dark_enabled):
        return gr.update(), None

    dark_mode_toggle.change(toggle_dark_mode, inputs=dark_mode_toggle, outputs=[])


    msg.submit(chat_with_bot, [msg, chatbot, state], [msg, chatbot, state])
    send_btn.click(chat_with_bot, [msg, chatbot, state], [msg, chatbot, state])

# === Launch ===
demo.launch()