################################################
#Pfad, wo Docs/Bilder/Filme abgelegt werden können - lokal, also hier im HF Space (sonst auf eigenem Rechner)
PATH_WORK = "."
CHROMA_DIR  = "/chroma/kkg"
CHROMA_PDF = './chroma/kkg/pdf'
CHROMA_WORD = './chroma/kkg/word'
CHROMA_EXCEL = './chroma/kkg/excel'
YOUTUBE_DIR = "/youtube"
HISTORY_PFAD = "/data/history"

###############################################
#URLs zu Dokumenten oder andere Inhalte, die einbezogen werden sollen
PDF_URL       = "https://arxiv.org/pdf/2303.08774.pdf"
WEB_URL       = "https://openai.com/research/gpt-4"
YOUTUBE_URL_1 = "https://www.youtube.com/watch?v=--khbXchTeE"
YOUTUBE_URL_2 = "https://www.youtube.com/watch?v=hdhZwyf24mE"
#YOUTUBE_URL_3 = "https://www.youtube.com/watch?v=vw-KWfKwvTQ"
#spezielle Webseiten als Datenbasis laden
urls = [
    "https://kkg.hamburg.de/unser-leitbild/"
    "https://kkg.hamburg.de/unsere-schulcharta/",
    "https://kkg.hamburg.de/koordination-unterrichtsentwicklung/",
    "https://kkg.hamburg.de/konzept-medien-und-it-am-kkg/",
]


##################################################
#Normalisierung eines Prompts
##################################################
def normalise_prompt (prompt):
    #alles Kleinbuchstaben
    prompt_klein =prompt.lower()
    #Word Tokenisation
    tokens = word_tokenize(prompt_klein)
    #Punktuierung entfernen
    tokens = [word for word in tokens if word.isalnum()]
    
    # Stop Word Entfernung
    nltk.download('stopwords')
    stop_words = set(stopwords.words('deutsch'))
    tokens = [word for word in tokens if not word in stop_words]
    # 5. Lemmatisierung: Worte in Grundform bringen, um Text besser vergleichen zu können
    nltk.download('wordnet')
    lemmatizer = WordNetLemmatizer()
    tokens = [lemmatizer.lemmatize(word) for word in tokens]
    # 6. Handling Special Characters (Remove non-alphanumeric characters)
    tokens = [re.sub(r'\W+', '', word) for word in tokens]    
    # 7. Spell Check (optional, using a library like pyspellchecker)
    from spellchecker import SpellChecker
    spell = SpellChecker()
    tokens = [spell.correction(word) for word in tokens]
    # Join tokens back to sentence
    normalized_prompt = ' '.join(tokens)
    print("normaiserd prompt..................................")
    print(normalized_prompt)
    return normalized_prompt


##################################################
#RAG Hilfsfunktionen - Dokumenten bearbeiten für Vektorstore
##################################################
##################################################
# Funktion, um für einen best. File-typ ein directory-loader zu definieren
def create_directory_loader(file_type, directory_path):
    #verscheidene Dokument loaders:
    loaders = {
        '.pdf': PyPDFLoader,
        '.word': UnstructuredWordDocumentLoader,
    }
    return DirectoryLoader(
        path=directory_path,
        glob=f"**/*{file_type}",
        loader_cls=loaders[file_type],
    )
################################################    
#die Inhalte splitten, um in Vektordatenbank entsprechend zu laden als Splits
def document_loading_splitting():
    ##############################
    # Document loading
    docs = []
    
    # kreiere einen DirectoryLoader für jeden file type
    pdf_loader = create_directory_loader('.pdf', CHROMA_PDF)
    word_loader = create_directory_loader('.word', CHROMA_WORD)
    print("PDF Loader done............................")
    
    # Load the files
    pdf_documents = pdf_loader.load()
    word_documents = word_loader.load()

    #alle zusammen in docs...
    docs.extend(pdf_documents)
    docs.extend(word_documents)

    #andere loader...
    # Load PDF
    #loader = PyPDFLoader(PDF_URL)
    #docs.extend(loader.load())
    # Load Web
    #loader = WebBaseLoader(WEB_URL)
    #docs.extend(loader.load())
    # Load YouTube
    #loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL_1,YOUTUBE_URL_2], PATH_WORK + YOUTUBE_DIR), OpenAIWhisperParser())
    #docs.extend(loader.load())
    ################################
    # Document splitting
    text_splitter = RecursiveCharacterTextSplitter(chunk_overlap = 150, chunk_size = 1500)
    splits = text_splitter.split_documents(docs)
    
    return splits

###########################################
#Chroma DB die splits ablegen - vektorisiert...
def document_storage_chroma(splits):
    #OpenAi embeddings----------------------------------
    Chroma.from_documents(documents = splits, embedding = OpenAIEmbeddings(disallowed_special = ()),  persist_directory = PATH_WORK + CHROMA_DIR)  

    #HF embeddings--------------------------------------
    #Chroma.from_documents(documents = splits, embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={"device": "cpu"}, encode_kwargs={'normalize_embeddings': False}),  persist_directory = PATH_WORK + CHROMA_DIR)  
    

############################################
#dokumente in chroma db vektorisiert ablegen können - die Db vorbereiten daüfur
def document_retrieval_chroma(llm, prompt):  
    #HF embeddings -----------------------------------
    #Alternative Embedding - für Vektorstore, um Ähnlichkeitsvektoren zu erzeugen - die ...InstructEmbedding ist sehr rechenaufwendig
    embeddings = HuggingFaceInstructEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": "cpu"})
    #etwas weniger rechenaufwendig:
    #embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={"device": "cpu"}, encode_kwargs={'normalize_embeddings': False})

    #ChromaDb um die embedings zu speichern
    db = Chroma(embedding_function = embeddings, persist_directory = PATH_WORK + CHROMA_DIR)
    return db


############################################
# rag_chain Alternative für RAg mit Bild-Upload, da hier das llm so nicht genutzt werden kann und der prompt mit den RAG Erweiterungen anders übergeben wird
#langchain nutzen, um prompt an llm zu leiten, aber vorher in der VektorDB suchen, um passende splits zum Prompt hinzuzufügen
#prompt mit RAG!!!
def rag_chain(prompt, db, k=3):
    rag_template = "Nutze ausschließlich die folgenden Kontext Teile am Ende, um die Frage zu beantworten . " + template + "Frage: " + prompt + "Kontext Teile: "
    retrieved_chunks  = db.similarity_search(prompt, k)

    neu_prompt = rag_template 
    for i, chunk in enumerate(retrieved_chunks):
        neu_prompt += f"{i+1}. {chunk}\n"

    return neu_prompt