#!/usr/bin/env python3
"""Полная индексация через sentence-transformers (CPU, стабильно)."""
import json, os, sys, time, glob, shutil
from chromadb import PersistentClient
from chromadb.utils import embedding_functions

ALL_DIRS = [
    "~/nubes/data/elsa_jsonl_wi_en",
    "~/nubes/data/elsa_mdb_jsonl",
    "~/nubes/data/elsa_mdb_jsonl/ipsvrap",
    "~/nubes/data/elsa_jsonl_htm_en",
    "~/nubes/data/elsa_jsonl",
    "~/nubes/data/elsa_mdb_jsonl/dbsvrfi",
    "~/nubes/data/elsa_mdb_jsonl/dbsvrfz",
]

DB_DIR = os.path.expanduser("~/nubes/chroma_db")
MODEL = "all-MiniLM-L6-v2"  # 90 MB, CPU

ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=MODEL)

if os.path.exists(DB_DIR):
    shutil.rmtree(DB_DIR)
    print("Old chroma_db deleted")

client = PersistentClient(path=DB_DIR)
collection = client.create_collection("elsa_docs", embedding_function=ef)

all_files = []
for d in ALL_DIRS:
    path = os.path.expanduser(d)
    files = sorted(glob.glob(os.path.join(path, "*.jsonl")))
    all_files.extend(files)
    print(f"{d}: {len(files)} files")

print(f"\nTotal: {len(all_files)} JSONL files")

total = 0
for fpath in all_files:
    docs, ids, metas = [], [], []
    prefix = os.path.basename(os.path.dirname(fpath)) + "_"
    
    try:
        with open(fpath) as fh:
            for line in fh:
                try:
                    d = json.loads(line)
                    title = d.get("title", "") or d.get("full_title", "")
                    text = d.get("text", "")
                    doc = (title + "\n" + text).strip()[:1000]
                    if len(doc) < 20:
                        continue
                    docs.append(doc)
                    ids.append(f"{prefix}{os.path.basename(fpath)}_{len(docs)}")
                    metas.append({"source": prefix.strip("_")})
                except:
                    continue
    except:
        continue
    
    if not docs:
        continue
    
    print(f"{os.path.basename(fpath)}: {len(docs)} docs", end="", flush=True)
    
    for j in range(0, len(docs), 500):
        batch = docs[j:j+500]
        batch_ids = ids[j:j+500]
        for retry in range(3):
            try:
                collection.add(ids=batch_ids, documents=batch,
                    metadatas=metas[j:j+500])
                total += len(batch)
                print(".", end="", flush=True)
                break
            except Exception as e:
                if retry < 2:
                    print(f"R{retry}", end="", flush=True)
                    time.sleep(5)
                else:
                    print(f"X", end="", flush=True)
    
    print(f" total={collection.count()}")

print(f"\n✅ Done. Total: {collection.count()} docs")