import chromadb 
import os
import hashlib

CHROMA_HOST = os.environ.get("CHROMA_HOST", "http://chroma.sirblob.co")
COLLECTION_NAME = "rag_documents"

_client =None 

def get_chroma_client ():
    global _client 
    if _client is None :
        _client =chromadb .HttpClient (host =CHROMA_HOST )
    return _client 

def get_collection (collection_name =COLLECTION_NAME ):
    client =get_chroma_client ()
    return client .get_or_create_collection (name =collection_name )

def insert_documents(texts, embeddings, collection_name=COLLECTION_NAME, metadata_list=None):
    collection = get_collection(collection_name)

    # Generate stable IDs using SHA256
    ids = [f"doc_{i}_{hashlib.sha256(text.encode('utf-8')).hexdigest()[:16]}" for i, text in enumerate(texts)]

    total_count = len(texts)
    batch_size = 2000
    
    for i in range(0, total_count, batch_size):
        end = min(i + batch_size, total_count)
        
        batch_ids = ids[i:end]
        batch_texts = texts[i:end]
        batch_embeddings = embeddings[i:end]
        batch_metadatas = metadata_list[i:end] if metadata_list else None
        
        if batch_metadatas:
            collection.add(
                ids=batch_ids,
                embeddings=batch_embeddings,
                documents=batch_texts,
                metadatas=batch_metadatas
            )
        else:
            collection.add(
                ids=batch_ids,
                embeddings=batch_embeddings,
                documents=batch_texts
            )

    return total_count

def search_documents (query_embedding ,collection_name =COLLECTION_NAME ,num_results =5 ,filter_metadata =None ):
    collection =get_collection (collection_name )

    query_params ={
    "query_embeddings":[query_embedding ],
    "n_results":num_results 
    }

    if filter_metadata :
        query_params ["where"]=filter_metadata 

    results =collection .query (**query_params )

    output =[]
    if results and results ["documents"]:
        for i ,doc in enumerate (results ["documents"][0 ]):
            score =results ["distances"][0 ][i ]if "distances"in results else None 
            meta =results ["metadatas"][0 ][i ]if "metadatas"in results else {}
            output .append ({
            "text":doc ,
            "score":score ,
            "metadata":meta 
            })

    return output 

def delete_documents_by_source (source_file ,collection_name =COLLECTION_NAME ):
    collection =get_collection (collection_name )
    results =collection .get (where ={"source":source_file })
    if results ["ids"]:
        collection .delete (ids =results ["ids"])
        return len (results ["ids"])
    return 0 

def get_all_metadatas (collection_name =COLLECTION_NAME ,limit =None ):
    collection =get_collection (collection_name )

    if limit :
        results =collection .get (include =["metadatas"],limit =limit )
    else :
        results =collection .get (include =["metadatas"])
    return results ["metadatas"]if results and "metadatas"in results else []
