import os
import uuid
import json
import chromadb
# from django.conf import settings


vector_db_path = "./speedbuild_vectordb"

default_collection = "default_collection"    

# def saveToVectorDB(collection_name,chunks,root_path,main_collection=False):
#     documents = []
#     metadatas = []

#     client = chromadb.PersistentClient(path="./speedbuild_vectordb")
#     collection = client.get_or_create_collection(name=collection_name)
    
#     if main_collection:
#         documents = [chunks['description']]
#         metadatas = [{
#             "template":chunks['template_name'],
#             "framework":chunks['framework'],
#         }]

#     else:

#         for chunk in chunks:
#             chunk_data = chunks[chunk]
#             documents.append(f"{chunk_data['doc']}\n\n{chunk_data['code']}" if chunk_data['doc'] is not None else chunk_data['code'])
#             metadatas.append({
#                 "source":os.path.relpath(chunk,root_path).split(":::")[0],
#                 "docs":chunk_data['doc'] if chunk_data['doc'] is not None else "Null",
#                 "dependencies": json.dumps(chunk_data.get('dependencies', []))
#             })

#     collection.add(
#         ids = [str(uuid.uuid4()) for _ in documents],
#         documents = documents,
#         metadatas = metadatas
#     )

def saveToVectorDB(doc,meta_data,collection_name=default_collection):

    client = chromadb.PersistentClient(path=vector_db_path)
    collection = client.get_or_create_collection(name=collection_name)
    doc_id = str(uuid.uuid4())

    collection.add(
        ids = [doc_id],
        documents = [doc],
        metadatas = [meta_data]
    )

    return doc_id

def query_collection(collection_name,query,framework=None,n=5):
    client = chromadb.PersistentClient(path=vector_db_path)
    collection = client.get_or_create_collection(name=collection_name)

    if framework:
        # filter collection meta data to match framework here
        response = collection.query(
            query_texts=query,
            n_results=n,
            where={"framework":framework}
        )
    else:
        response = collection.query(query_texts=query,n_results=n)

    context = []

    for i in range(0,len(response['documents'][0]),1):
        doc = response['documents'][0][i]
        meta_data = response['metadatas'][0][i]

        context.append({
            "name":meta_data['name'],
            "feature_id":meta_data['id'],
            "code":doc
        })

    return json.dumps(context)


def peek_collection(collection_name=default_collection):
    client = chromadb.PersistentClient(path=vector_db_path)
    collection = client.get_or_create_collection(name=collection_name)

    response = collection.peek()

    print(response)


def getCodeContext(query:str,framework:str):
    """
    Get Extra context on code base

    Args:
        query (str) : question you want to ask
        framework (str) : could be express or django

    Returns:
        code context
    """
    context = query_collection(default_collection,query,framework)
    return context


if __name__ == "__main__":
    context = getCodeContext("Pay fast payment processing and validation","django")
    print(context)