from ctransformers import AutoModelForCausalLM, AutoConfig
from sentence_transformers import SentenceTransformer
from chromadb.utils import embedding_functions
from chromadb.config import Settings
from pathlib import Path
import chromadb
import os
import json

# TheBloke/deepseek-coder-33B-instruct-GGUF  "ddh0/Yi-6B-200K-GGUF-fp16" 
#  "TheBloke/Mistral-7B-Code-16K-qlora-GGUF" # "TheBloke/Mistral-7B-Instruct-v0.1-GGUF" # "TheBloke/Mistral-7B-OpenOrca-GGUF" 
# "NousResearch/Yarn-Mistral-7b-128k" "JDWebProgrammer/custom_sft_adapter" 

MODEL_HF = "TheBloke/deepseek-coder-33B-instruct-GGUF"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

class AppModel:
    def __init__(self, embedding_model_name=EMBEDDING_MODEL, model=MODEL_HF, dataset_path="./data/logs", dir="./data", 
        context_limit=32000, temperature=0.8, max_new_tokens=4096, context_length=128000):
        self.model = model
        self.embedding_model_name = embedding_model_name
        self.model_config = AutoConfig.from_pretrained(self.model, context_length=context_length)
        self.emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=self.embedding_model_name.split("/")[1])
        self.chroma_client = chromadb.PersistentClient(path="./data/vectorstore", settings=Settings(anonymized_telemetry=False))
        self.sentences = [] 
        self.ref_collection = self.chroma_client.get_or_create_collection("ref", embedding_function=self.emb_fn)
        self.logs_collection = self.chroma_client.get_or_create_collection("logs", embedding_function=self.emb_fn)
        self.init_chroma()

        self.embedding_model = SentenceTransformer(self.embedding_model_name)
        self.llm = AutoModelForCausalLM.from_pretrained(self.model, model_type="mistral", config=self.model_config) #, cache_dir="./models" , gpu_layers=0 local_files_only=True) , 
        self.chat_log = []
        self.last_ai_response = ""
        self.last_user_prompt = "" 
        self.context_limit=context_limit
        self.temperature=temperature
        self.max_new_tokens=max_new_tokens

    def get_llm_query(self, input_prompt, user_prompt):
        self.last_user_prompt = str(user_prompt)
        new_response = self.llm(prompt=input_prompt, temperature=self.temperature, max_new_tokens=self.max_new_tokens) #, temperature=self.temperature, max_new_tokens=self.max_new_tokens)
        self.last_ai_response = str(new_response)
        self.save_file(f"[User_Prompt]: {user_prompt} \n[AI_Response]: {new_response} \n", "./data/logs/chat-log.txt")
        return new_response
        
    def get_embedding_values(self, input_str):
        tokenized_input = self.build_embeddings(input_str)
        print(tokenized_input)
        embedding_values = self.embedding_model.encode(tokenized_input)
        return embedding_values

    def get_embedding_docs(self, query_text, n_results=2):
        query_embeddings = self.get_embedding_values(query_text).tolist()[0]
        query_result = self.ref_collection.query(query_embeddings=query_embeddings,n_results=n_results) 
        return query_result["documents"]

    def init_chroma(self):
        docs, metas, ids = self.build_chroma_docs(directory="./data/reference", id_name="ref_")
        if docs:
            print(f"Loading Chroma (Reference) Docs: {len(docs)}")
            self.ref_collection.add(documents=docs, metadatas=metas, ids=ids)

        docs, metas, ids = self.build_chroma_docs(directory="./data/context", id_name="context_")
        if docs:
            print(f"Loading Chroma (Context) Docs: {len(docs)}")
            self.logs_collection.add(documents=docs, metadatas=metas, ids=ids)

    def build_chroma_docs(self, directory="./data/context", id_name="doc_", metatag={"source": "notion"}):
        directory = os.path.join(os.getcwd(), directory)
        docs = []
        metas = []
        ids = [] 
        fnum = 0
        for filename in os.listdir(directory):
            file_path = os.path.join(directory, filename)
            with open(file_path, 'r') as file:
                file_contents = file.read()
                splitter = "\n\n"
                if ".csv" in file_path:
                    splitter = "\n"
                anum = 0
                for a in file_contents.split(splitter): # split first by paragraph  
                    docs.append(a)
                    ids.append(id_name + str(fnum))
                    additional_metas = {"dir": directory, "filename":file_path, "chunk_number": anum }
                    metas.append({**metatag, **additional_metas})
                    fnum += 1
                    anum += 1
        docs = list(docs)
        metas = list(metas)
        ids = list(ids)
        return docs, metas, ids

    def build_embeddings(self, content, add_sentences=False):
        tokenized_sentences = []
        for b in content.split("\n"): # then by line 
            for c in b.split("  "): # then by tab
                for d in c.split(". "): # by sentence
                    tokenized_sentences.append(str(d))
                    if add_sentences:
                        self.sentences.append(str(d))
        return tokenized_sentences

    def save_file(self, data, filename="./data/context/chat-log.txt"):
        with open(filename, 'a') as f:
            f.write('\n\n' + str(data))
    
    def add_feedback(self, is_positive=True):
        feedback_str = ""
        if is_positive:
            feedback_str = "GOOD/PASS"
            self.chat_log.append(self.last_ai_response[:self.context_limit])
            self.save_file(self.last_ai_response)
        else:
            feedback_str = "BAD/FAIL"
        new_obj = f"[User_Prompt]: {self.last_user_prompt}\n[AI_Response]: {self.last_ai_response}\n[User_Feedback]: {feedback_str}\n\n"
        self.save_file(new_obj, "./data/logs/feedback-log.txt")


    def open_file(self, file_path):
        file_contents = ""
        with open(file_path, "r") as file:
            file_contents = file.read()
        return file_contents