import os
import tempfile
import pickle
from typing import List, Dict, Any

import streamlit as st
import numpy as np

from sentence_transformers import SentenceTransformer
import faiss
import PyPDF2
import docx
from groq import Groq

DEFAULT_EMBED_MODEL = "all-MiniLM-L6-v2"
DEFAULT_CHUNK_SIZE = 200
DEFAULT_CHUNK_OVERLAP = 50
DEFAULT_TOP_K = 5
INDEX_PATH = "faiss_index.bin"
METADATA_PATH = "metadata.pkl"

@st.cache_resource
def load_embedding_model(model_name: str = DEFAULT_EMBED_MODEL):
    return SentenceTransformer(model_name)


def save_uploaded_file(uploaded_file) -> str:
    tmpdir = tempfile.gettempdir()
    safe_name = os.path.basename(uploaded_file.name)
    temp_path = os.path.join(tmpdir, safe_name)
    with open(temp_path, "wb") as f:
        f.write(uploaded_file.getbuffer())
    return temp_path


def extract_text_from_pdf(path: str) -> str:
    text_pages = []
    try:
        reader = PyPDF2.PdfReader(path)
        for page in reader.pages:
            page_text = page.extract_text() or ""
            text_pages.append(page_text)
    except Exception:
        return ""
    return "\n\n".join(text_pages)


def extract_text_from_docx(path: str) -> str:
    try:
        doc = docx.Document(path)
        paragraphs = [p.text for p in doc.paragraphs]
        return "\n\n".join(paragraphs)
    except Exception:
        return ""


def extract_text_from_txt(path: str) -> str:
    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            return f.read()
    except Exception:
        return ""


def chunk_text(text: str, chunk_size: int = DEFAULT_CHUNK_SIZE, overlap: int = DEFAULT_CHUNK_OVERLAP) -> List[str]:
    words = text.split()
    chunks: List[str] = []
    start = 0
    n = len(words)
    while start < n:
        end = min(start + chunk_size, n)
        chunk = " ".join(words[start:end])
        chunks.append(chunk)
        if end == n:
            break
        start = max(0, end - overlap)
    return chunks


def build_faiss_index(embeddings: np.ndarray) -> faiss.IndexFlat:
    embeddings = np.ascontiguousarray(embeddings.astype('float32'))
    d = embeddings.shape[1]
    index = faiss.IndexFlatIP(d)
    index.add(embeddings)
    return index


def normalize_embeddings(vecs: np.ndarray) -> np.ndarray:
    vecs = np.asarray(vecs, dtype=np.float32)
    norms = np.linalg.norm(vecs, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return vecs / norms


def save_index_and_metadata(index: faiss.Index, metadata: List[Dict[str, Any]], index_path: str = INDEX_PATH, meta_path: str = METADATA_PATH):
    try:
        faiss.write_index(index, index_path)
        with open(meta_path, "wb") as f:
            pickle.dump(metadata, f)
    except Exception:
        pass


def load_index_and_metadata(index_path: str = INDEX_PATH, meta_path: str = METADATA_PATH):
    try:
        if os.path.exists(index_path) and os.path.exists(meta_path):
            index = faiss.read_index(index_path)
            with open(meta_path, "rb") as f:
                metadata = pickle.load(f)
            return index, metadata
    except Exception:
        return None, None
    return None, None


def get_groq_client_from_env() -> Groq | None:
    key = os.environ.get("GROQ_API_KEY") or os.environ.get("HF_GROQ_API_KEY") or os.environ.get("HF_API_TOKEN")
    if not key:
        return None
    return Groq(api_key=key)


def groq_chat_answer(client: Groq, messages: List[Dict[str, str]], model: str = "llama-3.3-70b-versatile", temperature: float = 0.0, max_tokens: int = 512) -> str:
    try:
        chat_completion = client.chat.completions.create(
            messages=messages,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        return getattr(chat_completion.choices[0].message, 'content', '')
    except Exception:
        return ""


def main():
    st.set_page_config(page_title="RAG with Groq + Sentence-Transformers", layout="wide")
    st.title("RAG (Retrieval Augmented Generation) — Groq + sentence-transformers + FAISS")

    st.sidebar.header("Index / Retrieval settings")
    chunk_size = int(st.sidebar.number_input("Chunk size (words)", min_value=50, max_value=2000, value=DEFAULT_CHUNK_SIZE, step=50))
    chunk_overlap = int(st.sidebar.number_input("Chunk overlap (words)", min_value=0, max_value=500, value=DEFAULT_CHUNK_OVERLAP, step=10))
    top_k = int(st.sidebar.number_input("Top-k results to retrieve", min_value=1, max_value=50, value=DEFAULT_TOP_K))

    st.sidebar.markdown("---")
    st.sidebar.write("Groq API key is loaded from Hugging Face secrets or environment variables. Do not paste keys in the UI.")

    with st.spinner("Loading embedding model..."):
        embed_model = load_embedding_model()

    index, metadata = load_index_and_metadata()
    if index is None and "index" in st.session_state:
        index = st.session_state["index"]
        metadata = st.session_state.get("metadata")

    if index is None:
        st.info("No existing FAISS index found. Upload documents and click 'Ingest documents' to build index.")
    else:
        try:
            nt = index.ntotal
        except Exception:
            nt = 0
        st.success(f"Loaded FAISS index with {nt} vectors.")

    st.header("1) Upload documents")
    uploaded_files = st.file_uploader("Upload PDF / DOCX / TXT files", type=["pdf", "docx", "txt"], accept_multiple_files=True)

    if uploaded_files:
        st.session_state["uploaded_list"] = [f.name for f in uploaded_files]
        st.write("Files ready to ingest:", st.session_state["uploaded_list"])

    if st.button("Ingest documents"):
        if not uploaded_files:
            st.warning("Please upload at least one file before ingesting.")
        else:
            all_chunks: List[str] = []
            metadata = []
            for uploaded in uploaded_files:
                tmp_path = save_uploaded_file(uploaded)
                name = uploaded.name
                text = ""
                if name.lower().endswith(".pdf"):
                    text = extract_text_from_pdf(tmp_path)
                elif name.lower().endswith(".docx"):
                    text = extract_text_from_docx(tmp_path)
                elif name.lower().endswith(".txt"):
                    text = extract_text_from_txt(tmp_path)
                else:
                    continue

                if not text or not text.strip():
                    continue

                chunks = chunk_text(text, chunk_size=chunk_size, overlap=chunk_overlap)
                for i, c in enumerate(chunks):
                    all_chunks.append(c)
                    metadata.append({"source": name, "chunk_id": i, "text": c})

            if not all_chunks:
                st.error("No chunks were generated. Aborting.")
            else:
                with st.spinner("Computing embeddings..."):
                    embeddings = embed_model.encode(all_chunks, show_progress_bar=True, convert_to_numpy=True)
                    embeddings = normalize_embeddings(embeddings)

                index = build_faiss_index(embeddings)
                st.session_state["index"] = index
                st.session_state["metadata"] = metadata
                try:
                    save_index_and_metadata(index, metadata)
                    st.success(f"Index built and saved. {len(all_chunks)} chunks indexed.")
                except Exception:
                    st.success(f"Index built in memory. {len(all_chunks)} chunks indexed.")
                st.info(f"Ingested {len(all_chunks)} chunks from {len(set(m['source'] for m in metadata))} file(s).")
                if len(all_chunks) > 0:
                    st.markdown("### Example chunk")
                    st.write(all_chunks[0][:1000])

    st.markdown("---")

    st.header("2) Ask questions (uses retrieved context + Groq generation)")
    query = st.text_area("Enter your question here", height=120)

    col1, col2 = st.columns([1, 3])
    with col1:
        if st.button("Ask"):
            if not query or not query.strip():
                st.warning("Please enter a question.")
            else:
                if "index" not in st.session_state or "metadata" not in st.session_state:
                    st.warning("No index available. Please ingest documents first.")
                else:
                    index = st.session_state["index"]
                    metadata = st.session_state["metadata"]
                    q_emb = embed_model.encode([query], convert_to_numpy=True)
                    q_emb = normalize_embeddings(q_emb)
                    q_emb = np.ascontiguousarray(q_emb.astype('float32'))
                    k = max(1, min(top_k, int(index.ntotal)))
                    if k == 0:
                        st.warning("Index is empty. Ingest documents first.")
                    else:
                        D, I = index.search(q_emb, k)
                        hits = I[0].tolist() if I is not None else []

                        contexts = []
                        for idx in hits:
                            if idx < 0 or idx >= len(metadata):
                                continue
                            meta = metadata[idx]
                            contexts.append(meta)

                        client = get_groq_client_from_env()
                        if client is None:
                            st.error("Groq client could not be created. Set GROQ_API_KEY in your Hugging Face secrets or environment.")
                        else:
                            system_msg = {
                                "role": "system",
                                "content": "You are an assistant that answers user questions using ONLY the context passages provided. If the answer is not contained in the context, say you don't know. Provide short answers and include the source(s) for any factual claims."
                            }

                            context_blocks = []
                            for c in contexts:
                                block = f"[Source: {c['source']} | chunk_id: {c['chunk_id']}]\n{c['text']}"
                                context_blocks.append(block)

                            context_text = "\n\n---\n\n".join(context_blocks)

                            user_msg_text = (
                                f"Context passages:\n\n{context_text}\n\n"
                                f"Question: {query}\n\n"
                                "Instructions: Use only the context passages above to answer. If the context does not contain enough information, respond that you don't know. Keep the answer concise and cite the source tags in square brackets."
                            )

                            user_msg = {"role": "user", "content": user_msg_text}

                            with st.spinner("Generating answer from Groq model..."):
                                answer = groq_chat_answer(client, messages=[system_msg, user_msg])

                            if answer:
                                st.subheader("Answer")
                                st.write(answer)

                                st.subheader("Retrieved passages (for transparency)")
                                for c in contexts:
                                    st.markdown(f"**Source:** {c['source']} — chunk {c['chunk_id']}")
                                    st.write(c['text'][:500])
                                    st.markdown("---")

    st.sidebar.markdown("---")
    st.sidebar.header("Index status")
    if "index" in st.session_state:
        try:
            st.sidebar.write(f"Vectors in index: {st.session_state['index'].ntotal}")
        except Exception:
            st.sidebar.write("Vectors in index: unknown")
    else:
        st.sidebar.write("No index built yet")

    st.sidebar.markdown("\nMade for development in Colab + Streamlit.\nUse responsibly and avoid publishing API keys.")


if __name__ == '__main__':
    main()