#!/usr/bin/env python3
"""Полная переиндексация ВСЕХ данных в ChromaDB. Работает до победного."""
import json, os, sys, time, urllib.request, glob
from chromadb import PersistentClient

EF_URL = "http://localhost:8081/embedding"
DB_DIR = os.path.expanduser("~/nubes/chroma_db")

ALL_DIRS = [
    "~/nubes/data/elsa_jsonl_wi_en",         # 244k WI XML
    "~/nubes/data/elsa_mdb_jsonl",            # 191k rldal
    "~/nubes/data/elsa_mdb_jsonl/ipsvrap",   # 2.5M parts
    "~/nubes/data/elsa_jsonl_htm_en",        # 7k hs2+www
    "~/nubes/data/elsa_jsonl",               # 24k HTM
    "~/nubes/data/elsa_mdb_jsonl/dbsvrfi",   # car ref
    "~/nubes/data/elsa_mdb_jsonl/dbsvrfz",   # PR codes
]

def embed(texts):
    for attempt in range(5):
        try:
            data = json.dumps({"content": texts}).encode()
            req = urllib.request.Request(EF_URL, data=data,
                headers={"Content-Type": "application/json"})
            with urllib.request.urlopen(req, timeout=120) as resp:
                body = json.loads(resp.read())
            if isinstance(body, list):
                return [item["embedding"][0] for item in body]
            return body.get("embedding", [])
        except Exception as e:
            print(f"\n  [embed retry {attempt+1}/5] {e}")
            time.sleep(15)
    raise Exception("Embedding failed after 5 retries")

def main():
    # Удалить старую БД и создать новую
    import shutil
    if os.path.exists(DB_DIR):
        shutil.rmtree(DB_DIR)
        print("Старая ChromaDB удалена")
    
    client = PersistentClient(path=DB_DIR)
    collection = client.create_collection("elsa_docs")
    
    # Собираем все JSONL
    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"\nВсего {len(all_files)} JSONL файлов")
    
    total = 0
    errors = 0
    
    for fpath in all_files:
        docs, ids = [], []
        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)}")
                    except:
                        continue
        except Exception as e:
            print(f"  [SKIP {os.path.basename(fpath)}] read error: {e}")
            errors += 1
            continue
        
        if not docs:
            print(f"  {os.path.basename(fpath)}: empty")
            continue
        
        print(f"  {os.path.basename(fpath)}: {len(docs)} docs", end="", flush=True)
        
        # Бачами по 50
        for j in range(0, len(docs), 50):
            batch_docs = docs[j:j+50]
            batch_ids = ids[j:j+50]
            
            for retry in range(5):
                try:
                    embs = embed(batch_docs)
                    collection.add(ids=batch_ids, documents=batch_docs,
                        embeddings=embs, metadatas=[{"source": prefix.strip("_")}] * len(batch_docs))
                    total += len(batch_docs)
                    print(".", end="", flush=True)
                    break
                except Exception as e:
                    if retry < 4:
                        print(f"R{retry}", end="", flush=True)
                        time.sleep(20)
                    else:
                        print(f"X", end="", flush=True)
                        errors += len(batch_docs)
            
            time.sleep(0.2)
        
        print(f" total={collection.count()}")
    
    print(f"\n✅ Done. Total: {collection.count()} docs. Errors: {errors}")

if __name__ == "__main__":
    main()