from dotenv import load_dotenv
from google import genai
import streamlit as st
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
import chromadb
import os
from chromadb.utils import embedding_functions
from huggingface_hub import InferenceClient

# ------------------------
# Setup
# ------------------------
st.set_page_config(page_title="Simple RAG", layout="wide")

# Load the environment variables from the .env file
load_dotenv()

# Initialize session state with API keys
if "gemini_key" not in st.session_state:
    st.session_state.gemini_key = os.getenv("GEMINI_API_KEY", "")

st.title("📄 DocuChat")

model = SentenceTransformer("all-MiniLM-L6-v2")

# Show forms only if keys are missing
if not st.session_state.gemini_key:
    st.warning("⚠️ Gemini API Key is not set")
    gemini_input = st.text_input("Enter your Gemini API Key", type="password")
    if gemini_input:
        st.session_state.gemini_key = gemini_input
        st.rerun()

GEMINI_API_KEY = st.session_state.gemini_key

if not GEMINI_API_KEY:
    st.error("❌ Please provide your Gemini API Key to continue")
    st.stop()

genai_client = genai.Client(api_key=GEMINI_API_KEY)

chroma_client = chromadb.Client()
# Delete and recreate collection to ensure consistent embedding dimensions
try:
    chroma_client.delete_collection(name="docs")
except:
    pass
collection = chroma_client.get_or_create_collection(name="docs")

# ------------------------
# Helpers
# ------------------------
def extract_text(file):
    if file.type == "application/pdf":
        reader = PdfReader(file)
        text = ""
        for page in reader.pages:
            text += page.extract_text() or ""
        return text
    else:
        return file.read().decode("utf-8")


def chunk_text(text, chunk_size=500, overlap=50):
    chunks = []
    start = 0

    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap

    return chunks

def query_embedding_api(texts):
    """Helper function to fetch embeddings using SentenceTransformer."""
    try:
        embeddings = model.encode(texts)
        # Convert to list if needed
        if hasattr(embeddings, 'tolist'):
            embeddings = embeddings.tolist()
        return embeddings
    except Exception as e:
        st.error(f"Error generating embeddings: {e}")
        return None

def embed_and_store(chunks):
    # embeddings = model.encode(chunks).tolist()
    embeddings = query_embedding_api(chunks)
    
    if embeddings is None:
        st.error("Failed to generate embeddings")
        return

    for i, chunk in enumerate(chunks):
        collection.add(
            documents=[chunk],
            embeddings=[embeddings[i]], 
            ids=[f"id_{i}"]
        )


def retrieve(query, k=5):
    # Fetch embedding for a single query string
    api_response = query_embedding_api([query])
    
    if api_response is None:
        return []
    
    query_embedding = api_response[0]  

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=k
    )

    return results["documents"][0]


# ------------------------
# UI
# ------------------------
tab1, tab2 = st.tabs(["Upload", "Chat"])

# ------------------------
# Upload Page
# ------------------------
with tab1:
    st.header("Upload Document")

    file = st.file_uploader("Upload PDF or TXT", type=["pdf", "txt"])

    if file:
        text = extract_text(file)
        chunks = chunk_text(text)
        
        embed_and_store(chunks)

        st.success(f"Stored {len(chunks)} chunks!")

# ------------------------
# Chat Page
# ------------------------
with tab2:
    st.header("Ask Questions")

    query = st.text_input("Enter your question")

    if query:
        docs = retrieve(query)

        context = "\n\n".join(docs)

        # 2. Construct a clean system instruction and prompt for the LLM
        prompt = f"""
        You are a helpful assistant. Answer the question based ONLY on the provided context below.
        If the answer cannot be found in the context, say "I cannot find the answer in the provided documents."

        Context:
        {context}

        Question:
        {query}
        """
        
        # 3. Call Gemini using the smallest, high-speed model
        with st.spinner("Thinking..."):
            try:
                response = genai_client.models.generate_content(
                    model='gemini-2.5-flash',
                    contents=prompt,
                )
                
                st.markdown("## Response")
                st.markdown(response.text)
                
            except Exception as e:
                st.error(f"Gemini API Error: {e}")