import os
import chromadb
from chromadb.utils import embedding_functions
from django.conf import settings

# Tetapkan lokasi folder pangkalan data vektor
CHROMA_PATH = os.path.join(settings.BASE_DIR, "chroma_db")

# Ganti pemulaan global dengan pemuatan malas (Lazy Initialization)
_chroma_client = None
_embedding_model = None

def get_chroma_client():
    """Dapatkan client ChromaDB hanya apabila dipanggil (Mencegah Deadlock)"""
    global _chroma_client
    if _chroma_client is None:
        _chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
    return _chroma_client

def get_embedding_model():
    """Dapatkan enjin pembenaman hanya apabila diperlukan"""
    global _embedding_model
    if _embedding_model is None:
        # Dipaksa guna model lokal secara mutlak bagi mengelakkan ralat 403 Forbidden dari OpenAI
        print("🚀 [ChromaDB] Menggunakan model lokal (SentenceTransformer) secara mutlak.")
        _embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
    return _embedding_model

def get_bengkel_collection(bengkel_id):
    """Cipta atau dapatkan 'rak buku' (collection) khas untuk satu bengkel."""
    client = get_chroma_client()
    model = get_embedding_model()
    
    collection_name = f"bengkel_rujukan_{bengkel_id}"
    collection = client.get_or_create_collection(
        name=collection_name,
        embedding_function=model
    )
    return collection

def chunk_text(text, chunk_size=400, overlap=50):
    """Cincang teks panjang kepada ketulan (chunks)."""
    words = text.split()
    chunks = []
    if not words:
        return chunks
    for i in range(0, len(words), chunk_size - overlap):
        chunk = " ".join(words[i:i + chunk_size])
        chunks.append(chunk)
    return chunks

def simpan_dokumen_ke_vektor(bengkel_id, dokumen_id, nama_dokumen, teks_penuh):
    """Fungsi utama untuk potong teks dan masukkan ke dalam ChromaDB."""
    if not teks_penuh.strip():
        return

    collection = get_bengkel_collection(bengkel_id)
    chunks = chunk_text(teks_penuh)
    
    ids = []
    documents = []
    metadatas = []
    
    for i, chunk in enumerate(chunks):
        ids.append(f"doc_{dokumen_id}_chunk_{i}")
        documents.append(chunk)
        metadatas.append({
            "dokumen_id": dokumen_id,
            "nama_dokumen": nama_dokumen,
            "chunk_index": i
        })
    
    collection.upsert(
        ids=ids,
        documents=documents,
        metadatas=metadatas
    )
    print(f"✅ Berjaya simpan {len(chunks)} chunks vektor untuk fail: {nama_dokumen}")

def cari_maklumat_berkaitan(bengkel_id, soalan, jumlah_jawapan=3):
    """Mencari 'chunk' paling tepat dalam ChromaDB berdasarkan soalan."""
    collection = get_bengkel_collection(bengkel_id)
    results = collection.query(
        query_texts=[soalan],
        n_results=jumlah_jawapan
    )
    
    if not results["documents"] or not results["documents"][0]:
        return ""
    
    teks_dijumpai = "\n...\n".join(results["documents"][0])
    return teks_dijumpai