import json
import os
import re
from pathlib import Path
import streamlit as st
from rag.ingestion_pipeline import process_doc, create_vector_store
from rag.retrieval_pipeline import retrieval_pipeline, load_vector_store, chat_history
import hashlib

DB_DIR = Path("db")
INDEX_PATH = DB_DIR / "index.json" 

st.title("RAG Assistant")


def load_index():
    """ Load the stored filename and his path. """

    if INDEX_PATH.exists():
        return json.loads(INDEX_PATH.read_text())
    return {}


def save_index(data):
    """ Save the filename mapping to disk """

    # Persist the filename -> db_path mapping on disk.
    DB_DIR.mkdir(parents=True, exist_ok=True)
    INDEX_PATH.write_text(json.dumps(data, indent=2))

def file_hash(uploaded_file) -> str :
    """ Return the content of the file uploaded in a short MD5 hash """

    hash_md5 = hashlib.md5()
    hash_md5.update(uploaded_file.getbuffer())
    return hash_md5.hexdigest()[:10] 

def safe_name(file_name: str, uploaded_file) -> str:
    """ Convert The uploaded filename into a unique, safe and clean folder name"""

    # Convert the uploaded filename into a safe folder name.
    # This avoids spaces and special characters in the DB path.
    clean_name = re.sub(r"[^a-zA-Z0-9_-]", "_", Path(file_name).stem)
    file_name_hashed = file_hash(uploaded_file)
    return f"{clean_name}_{file_name_hashed}"


def get_db_path(filename: str, uploaded_file) -> str:
    """ Return the local Chroma database path for a specific uploaded name."""

    # Each filename gets its own Chroma folder.
    # This avoids deleting an open DB and causing readonly/locked errors.
    return str(DB_DIR / f"chroma_db_{safe_name(filename, uploaded_file)}")


# Load the filename registry once at startup.
if "file_index" not in st.session_state:
    st.session_state.file_index = load_index()

if "vector_store" not in st.session_state:
    st.session_state.vector_store = None
    st.session_state.current_document = None

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

# If a different file is selected, clear the current in-memory store.
if uploaded and uploaded.name != st.session_state.current_document:
    st.info("A new file was selected. Process it to load or create its document store.")
    st.session_state.vector_store = None

# Process only when the button is clicked.
if uploaded and st.button("Process Document"):
    file_name = uploaded.name
    registry_key = safe_name(file_name, uploaded)
    db_path = st.session_state.file_index.get(registry_key, get_db_path(file_name, uploaded))

    # If this file was already processed before, reuse its existing vector store.
    if os.path.exists(db_path):
        st.session_state.vector_store = load_vector_store(db_path)
        st.session_state.current_document = file_name
        st.success("Existing vector store loaded for this file.")
    else:
        # New file: parse, chunk, and store it in a new Chroma folder.
        chunked_text = process_doc(uploaded, file_hash(uploaded))

        if not chunked_text:
            st.warning("No text chunks were created from the uploaded file.")
        else:
            # Clear chat history because this is a different document.
            chat_history.clear()

            with st.spinner("Creating vector store..."):
                st.session_state.vector_store = create_vector_store(
                    chunked_text,
                    path=db_path,
                )

            # Save the filename -> db_path mapping so we can reuse it later.
            st.session_state.file_index[registry_key] = db_path
            save_index(st.session_state.file_index)

            st.session_state.current_document = file_name
            st.success(f"Split into {len(chunked_text)} chunks")
            st.success("Document processed and saved.")

if st.session_state.vector_store is None:
    st.info("Upload and process a document before asking questions.")

question = st.text_input(
    "Ask a question about your document",
    disabled=st.session_state.vector_store is None,
)

if question and st.session_state.vector_store:
    answer = retrieval_pipeline(question, st.session_state.vector_store)
    st.write(answer)