import os
from .connection import get_mongo_client

DB_NAME = "ethix_vectors"
# Default collection name to match Chroma's default
DEFAULT_COLLECTION_NAME = "rag_documents"

def get_collection(collection_name=DEFAULT_COLLECTION_NAME):
    client = get_mongo_client()
    db = client[DB_NAME]
    return db[collection_name]

def insert_documents(texts, embeddings, collection_name=DEFAULT_COLLECTION_NAME, metadata_list=None):
    collection = get_collection(collection_name)
    
    docs = []
    for i, text in enumerate(texts):
        doc = {
            "text": text,
            "embedding": embeddings[i],
            "metadata": metadata_list[i] if metadata_list else {}
        }
        # Flatten metadata for easier filtering if needed, or keep in 'metadata' subdoc
        # Atlas Vector Search usually suggests keeping it cleaner, but 'metadata' field is fine
        # provided index is on metadata.field
        
        # Add source file to top level for easier deletion
        if metadata_list and "source" in metadata_list[i]:
             doc["source"] = metadata_list[i]["source"]
             
        docs.append(doc)
    
    if docs:
        result = collection.insert_many(docs)
        return len(result.inserted_ids)
    return 0

def search_documents(query_embedding, collection_name=DEFAULT_COLLECTION_NAME, num_results=5, filter_metadata=None):
    collection = get_collection(collection_name)
    
    # Construct Atlas Search Aggregation
    # Note: 'index' name depends on what user created. Default is often 'default' or 'vector_index'.
    # We'll use 'vector_index' as a convention.
    
    # Filter construction
    pre_filter = None
    if filter_metadata:
        # Atlas Search filter syntax is different from simple match
        # It requires MQL (Mongo Query Language) match or compound operators
        # For simplicity in $vectorSearch, strict filters are passed in 'filter' field (if serverless/v2)
        # or separate match stage.
        
        # Modern Atlas Vector Search ($vectorSearch) supports 'filter' inside the operator
        
        must_clauses = []
        for k, v in filter_metadata.items():
            # Assuming filter_metadata matches metadata fields inside 'metadata' object
            # or we map specific known fields.
            # In Chroma, filters are flat. In insert above, we put them in 'metadata'.
            must_clauses.append({
                "term": {
                    "path": f"metadata.{k}",
                    "value": v
                }
            })
            
        if must_clauses:
            if len(must_clauses) == 1:
                pre_filter = must_clauses[0]
            else:
                pre_filter = {
                    "compound": {
                        "filter": must_clauses
                    }
                }

    pipeline = [
        {
            "$vectorSearch": {
                "index": "vector_index",
                "path": "embedding",
                "queryVector": query_embedding,
                "numCandidates": num_results * 10,
                "limit": num_results
            }
        }
    ]
    
    # Apply filter if exists
    if pre_filter:
        pipeline[0]["$vectorSearch"]["filter"] = pre_filter

    pipeline.append({
        "$project": {
            "_id": 0,
            "text": 1,
            "metadata": 1,
            "score": {"$meta": "vectorSearchScore"}
        }
    })

    try:
        results = list(collection.aggregate(pipeline))
        
        # Format to match Chroma output style
        # Chroma: [{'text': ..., 'score': ..., 'metadata': ...}]
        output = []
        for doc in results:
            output.append({
                "text": doc.get("text", ""),
                "score": doc.get("score", 0),
                "metadata": doc.get("metadata", {})
            })
        return output
        
    except Exception as e:
        print(f"Atlas Vector Search Error: {e}")
        # Fallback to empty if index doesn't exist etc
        return []

def delete_documents_by_source(source_file, collection_name=DEFAULT_COLLECTION_NAME):
    collection = get_collection(collection_name)
    # We added 'source' to top level in insert_documents for exactly this efficiency
    result = collection.delete_many({"source": source_file})
    return result.deleted_count

def get_all_metadatas(collection_name=DEFAULT_COLLECTION_NAME, limit=None):
    collection = get_collection(collection_name)
    
    cursor = collection.find({}, {"metadata": 1, "source": 1, "_id": 0})
    if limit:
        cursor = cursor.limit(limit)
        
    metadatas = []
    for doc in cursor:
        meta = doc.get("metadata", {})
        # Ensure 'source' is in metadata if we pulled it out
        if "source" in doc and "source" not in meta:
            meta["source"] = doc["source"]
        metadatas.append(meta)
        
    return metadatas
