import os
import zipfile
import chromadb


# =========================
# 1. FIND ZIP FILES
# =========================
zip_files = [f for f in os.listdir(".") if f.endswith(".zip")]

print("ZIP FILES:", zip_files)


# =========================
# 2. EXTRACT
# =========================
for z in zip_files:
    try:
        with zipfile.ZipFile(z, "r") as zip_ref:
            zip_ref.extractall(z.replace(".zip", ""))
        print("✔ Extracted:", z)
    except Exception as e:
        print("ZIP ERROR:", z, e)


# =========================
# 3. FIND CHROMA ROOT (IMPORTANT FIX)
# =========================
def find_chroma_root():
    for root, dirs, files in os.walk("."):
        if "chroma.sqlite3" in files:
            # IMPORTANT: return root directory (not file level)
            return root
    return None


db_path = find_chroma_root()

if not db_path:
    print("❌ No Chroma DB found!")
    exit()

print("🔥 CHROMA ROOT:", db_path)


# =========================
# 4. LOAD CHROMA SAFELY
# =========================
try:
    client = chromadb.PersistentClient(
        path=db_path,
        settings=chromadb.Settings(
            anonymized_telemetry=False
        )
    )
except Exception as e:
    print("❌ CLIENT ERROR:", e)
    exit()


# =========================
# 5. CHECK COLLECTIONS
# =========================
try:
    collections = client.list_collections()
except Exception as e:
    print("❌ LIST COLLECTION ERROR:", e)
    exit()

print("\n===== COLLECTIONS =====")

if not collections:
    print("❌ No collections found!")
    exit()

for c in collections:
    col = client.get_collection(c.name)
    print(f"✔ {c.name} | COUNT = {col.count()}")


# =========================
# 6. TEST RETRIEVAL
# =========================
def test(query="leadership"):

    print("\n===== TEST QUERY =====")

    for c in collections:

        col = client.get_collection(c.name)

        try:
            res = col.query(
                query_texts=[query],
                n_results=5
            )

            docs = res.get("documents", [[]])[0]

            print(f"\n--- {c.name} ---")
            print("docs:", len(docs))

            for d in docs:
                print("-", d[:120])

        except Exception as e:
            print("QUERY ERROR:", e)


test()