# app.py
import os
import time
import traceback
from typing import List

import streamlit as st
from PIL import Image
import requests

from utils import (
    extract_text_from_pdf,
    crawl_site_dynamic,
    chunk_text,
    compute_embeddings,
    retrieve_top_k,
    save_embeddings,
    load_embeddings,
)

# ---------- Config ----------
EMBED_FILE = "embeddings.pkl"
DEFAULT_PDF = "Prospectus-Dawood-University-2024-2025-Final-Version-2 (1).pdf"
LOGO_FILE = "duet_logo.png"
CRAWL_START = "https://duet.edu.pk/"
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
HF_GEN_MODEL = "google/flan-t5-large"

st.set_page_config(page_title="🧠 DUET Chatbot", layout="wide")

# ---------- UI CSS ----------
st.markdown("""
<style>
.topbar { background: linear-gradient(90deg,#0b84ff,#6f0dd3); color:white; padding:16px; border-radius:12px; }
.title { font-size:26px; font-weight:700; }
.subtitle { color:#e6f0ff; margin-top:4px; }
.chat-window { max-width:100%; }
.msg-user { background:#e6fffa; border-radius:10px; padding:12px; margin:6px 0; color:#064e3b; }
.msg-bot { background:#f1f5f9; border-radius:10px; padding:12px; margin:6px 0; color:#0f172a; }
.small { font-size:12px; color: #6b7280; }
</style>
""", unsafe_allow_html=True)

# ---------- Header ----------
cols = st.columns([1, 8, 1])
with cols[1]:
    if os.path.exists(LOGO_FILE):
        try:
            st.image(Image.open(LOGO_FILE), width=110)
        except Exception:
            pass
    st.markdown('<div class="topbar"><div class="title">🧠 DUET Chatbot</div>'
                '<div class="subtitle">Prospectus + DUET website powered assistant — answers use official content only.</div></div>',
                unsafe_allow_html=True)

# ---------- Session state ----------
if "chunks" not in st.session_state:
    st.session_state["chunks"] = []
if "vectors" not in st.session_state:
    st.session_state["vectors"] = None
if "sources" not in st.session_state:
    st.session_state["sources"] = []
if "index_ready" not in st.session_state:
    st.session_state["index_ready"] = False
if "history" not in st.session_state:
    st.session_state["history"] = []  # list of (q,a)
if "last_query_time" not in st.session_state:
    st.session_state["last_query_time"] = 0.0

# ---------- Helper: LLM calls ----------
def call_groq_chat(system_prompt: str, context: str, question: str):
    if not GROQ_API_KEY:
        return None, "Groq key not configured"
    url = "https://api.groq.com/openai/v1/chat/completions"
    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
    payload = {
        "model":"llama-3.3-70b-versatile",
        "messages":[
            {"role":"system","content": system_prompt},
            {"role":"user","content": f"Use ONLY the context below to answer the question. If not present say you couldn't find it.\n\nContext:\n{context}\n\nQuestion: {question}"}
        ],
        "temperature": 0.2,
        "max_tokens": 512
    }
    try:
        r = requests.post(url, headers=headers, json=payload, timeout=60)
        if r.status_code == 200:
            return r.json()["choices"][0]["message"]["content"], None
        return None, f"Groq returned {r.status_code}: {r.text}"
    except Exception as e:
        return None, str(e)

def call_hf_generate(prompt: str):
    if not HF_API_TOKEN:
        return None, "HF token not set"
    url = f"https://api-inference.huggingface.co/models/{HF_GEN_MODEL}"
    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 256, "temperature": 0.2}}
    try:
        r = requests.post(url, headers=headers, json=payload, timeout=60)
        if r.status_code == 200:
            jr = r.json()
            if isinstance(jr, list) and "generated_text" in jr[0]:
                return jr[0]["generated_text"], None
            if isinstance(jr, dict) and "generated_text" in jr:
                return jr["generated_text"], None
            return str(jr), None
        return None, f"HF returned {r.status_code}: {r.text}"
    except Exception as e:
        return None, str(e)

# ---------- Build or load embeddings (preprocess) ----------
def ensure_embeddings_ready(force_rebuild: bool = False):
    """
    - If embeddings.pkl exists and force_rebuild is False -> load it
    - Otherwise: extract PDF text, crawl site, chunk, embed, save to embeddings.pkl
    """
    if st.session_state["index_ready"] and not force_rebuild:
        return True

    # If file exists, load
    if os.path.exists(EMBED_FILE) and not force_rebuild:
        try:
            chunks, vectors, sources = load_embeddings_from_file()
            if chunks and vectors is not None:
                st.session_state["chunks"] = chunks
                st.session_state["vectors"] = vectors
                st.session_state["sources"] = sources
                st.session_state["index_ready"] = True
                return True
        except Exception:
            # continue to rebuild if load fails
            pass

    # Build pipeline (slow step) but only run once
    st.info("Building knowledge index from prospectus + DUET site. This runs once and may take ~30-90s on CPU.")
    # PDF
    pdf_text = ""
    if os.path.exists(DEFAULT_PDF):
        try:
            pdf_text = extract_text_from_pdf(DEFAULT_PDF)
        except Exception as e:
            st.warning(f"PDF extraction error: {e}")
    else:
        st.warning("Default prospectus not found in repo root; PDF extraction skipped.")

    # Crawl site
    try:
        site_text, visited_urls = crawl_site_dynamic(CRAWL_START, max_pages=120, max_depth=2)
    except Exception as e:
        site_text = ""
        visited_urls = []
        st.warning(f"Site crawling failed: {e}")

    combined = (pdf_text or "") + "\n\n" + (site_text or "")
    if not combined.strip():
        st.error("No content available from prospectus or site to build index.")
        return False

    # chunk
    chunks = chunk_text(combined, chunk_size_words=220, overlap=40)
    # embed
    vectors = compute_embeddings(chunks)

    # save
    try:
        save_embeddings_to_file(chunks, vectors, sources=visited_urls)
    except Exception:
        st.warning("Could not save embeddings to file; proceeding with in-memory data only.")

    st.session_state["chunks"] = chunks
    st.session_state["vectors"] = vectors
    st.session_state["sources"] = visited_urls
    st.session_state["index_ready"] = True
    st.success(f"Index ready with {len(chunks)} chunks.")
    return True

# wrappers to save/load via utils (to keep names stable)
def save_embeddings_to_file(chunks, vectors, sources):
    save_embeddings(EMBED_FILE, chunks, vectors, sources)

def load_embeddings_from_file():
    from utils import load_embeddings
    return load_embeddings(EMBED_FILE)

# ---------- Main chat logic ----------
SYSTEM_PROMPT = (
    "You are a DUET assistant. Use ONLY the provided context (extracted from the prospectus and DUET website) to answer "
    "the user's question. If the information isn't in the provided context, explicitly reply: "
    "'I could not find this information in the official university material.' Be concise and student-friendly."
)

# Input area (single page chat)
st.markdown("### 💬 Chat with DUET")
input_col, btn_col = st.columns([9,1])
with input_col:
    question = st.text_input("Ask anything about DUET...", key="q")
with btn_col:
    pressed = st.button("Send")

# When user presses send
if pressed and question:
    start_time = time.time()
    # Ensure embeddings ready (load or build once)
    ok = ensure_embeddings_ready()
    if not ok:
        st.error("Index build failed. Please check logs.")
    else:
        # retrieval
        with st.spinner("Retrieving relevant passages..."):
            chunks = st.session_state["chunks"]
            vectors = st.session_state["vectors"]
            retrieved, scores, indices = retrieve_top_k(question, chunks, vectors, top_k=5)

        if not retrieved:
            answer = "I could not find this information in the official university material."
        else:
            # craft context for LLM
            context = "\n\n---\n\n".join(retrieved)
            answer = None
            # prefer Groq
            if GROQ_API_KEY:
                with st.spinner("Generating polished answer via Groq..."):
                    try:
                        resp, err = call_groq_chat(SYSTEM_PROMPT, context, question)
                        if resp:
                            answer = resp.strip()
                    except Exception as e:
                        answer = None
            # fallback HF
            if answer is None and HF_API_TOKEN:
                with st.spinner("Generating polished answer via Hugging Face..."):
                    try:
                        prompt = f"{SYSTEM_PROMPT}\n\nContext:\n{context}\n\nQuestion: {question}\n\nAnswer:"
                        resp, err = call_hf_generate(prompt)
                        if resp:
                            answer = resp.strip()
                    except Exception:
                        answer = None
            # final fallback: extractive summary
            if answer is None:
                # combine top retrieved chunks into extractive answer
                pieces = []
                for r in retrieved[:3]:
                    pieces.append(r.strip()[:900])
                answer = "\n\n".join(pieces)

        # Save history and show with typewriter animation in a placeholder
        st.session_state["history"].append((question, ""))  # placeholder
        placeholder = st.empty()
        # Render all previous exchanges above placeholder
        def render_conversation():
            for q,a in st.session_state["history"][:-1]:
                st.markdown(f"<div class='msg-user'><strong>You:</strong> {q}</div>", unsafe_allow_html=True)
                st.markdown(f"<div class='msg-bot'><strong>DUET Bot:</strong> {a}</div>", unsafe_allow_html=True)

        render_conversation()

        # Typewriter animate the final answer in placeholder
        typed = ""
        for ch in answer:
            typed += ch
            # update placeholder content
            placeholder.markdown(f"<div class='msg-user'><strong>You:</strong> {question}</div>"
                                 f"<div class='msg-bot'><strong>DUET Bot:</strong> {typed}</div>",
                                 unsafe_allow_html=True)
            time.sleep(0.006)  # adjust speed as desired

        # finalize history
        st.session_state["history"][-1] = (question, answer)

        # show evidence under collapsible
        with st.expander("📚 Retrieved evidence (click to expand)"):
            for i, chunk in enumerate(retrieved):
                score = scores[i] if i < len(scores) else None
                idx = indices[i] if i < len(indices) else None
                st.markdown(f"**Passage {i+1}** (score {score:.3f}, chunk index {idx}):")
                snippet = chunk[:1500] + ("..." if len(chunk) > 1500 else "")
                st.write(snippet)

        elapsed = time.time() - start_time
        st.markdown(f"<div class='small'>Response time: {elapsed:.2f}s</div>", unsafe_allow_html=True)

# Render full conversation if no new message (page load)
if not pressed:
    for q,a in st.session_state["history"]:
        st.markdown(f"<div class='msg-user'><strong>You:</strong> {q}</div>", unsafe_allow_html=True)
        st.markdown(f"<div class='msg-bot'><strong>DUET Bot:</strong> {a}</div>", unsafe_allow_html=True)

# ---------- Footer ----------
st.markdown("<hr><small>Powered by prospectus & https://duet.edu.pk/ — if critical, cross-check official DUET sources.</small>", unsafe_allow_html=True)