#!/usr/bin/env python3
"""Ingest WI JSONL into ChromaDB — one file at a time, batches of 50."""
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")
# Default: ingest ALL JSONL dirs
ALL_DIRS = [
    os.path.expanduser("~/nubes/data/elsa_jsonl_wi_en"),        # 244k WI XML
    os.path.expanduser("~/nubes/data/elsa_mdb_jsonl"),            # 191k rldal
    os.path.expanduser("~/nubes/data/elsa_mdb_jsonl/ipsvrap"),   # 2.5M parts
    os.path.expanduser("~/nubes/data/elsa_jsonl_htm_en"),        # 7k hs2+www
    os.path.expanduser("~/nubes/data/elsa_jsonl"),               # 24k HTM
    os.path.expanduser("~/nubes/data/elsa_mdb_jsonl/dbsvrfi"),   # car ref
    os.path.expanduser("~/nubes/data/elsa_mdb_jsonl/dbsvrfz"),   # PR codes
]

DATA_DIR = sys.argv[1] if len(sys.argv) > 1 else None
COLLECTION_NAME = "elsa_docs"

# Connect to existing ChromaDB
client = PersistentClient(path=DB_DIR)
try:
    collection = client.get_or_create_collection(COLLECTION_NAME)
except:
    collection = client.create_collection(COLLECTION_NAME)

existing = collection.count()
print(f"Collection has {existing} docs already")

# Find remaining files
if DATA_DIR:
    files = sorted(glob.glob(os.path.join(DATA_DIR, "*.jsonl")))
else:
    files = sorted(glob.glob(os.path.join(ALL_DIRS[0], "*.jsonl")))
    for d in ALL_DIRS[1:]:
        files.extend(sorted(glob.glob(os.path.join(d, "*.jsonl"))))

print(f"Processing {len(files)} JSONL files")

for fpath in files:
    docs, ids = [], []
    prefix = os.path.basename(os.path.dirname(fpath)) + "_"
    with open(fpath) as fh:
        for line in fh:
            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)}")
    
    if not docs:
        print(f"{os.path.basename(fpath)}: empty, skip")
        continue
    
    print(f"{os.path.basename(fpath)}: {len(docs)} docs", end="")
    
    for j in range(0, len(docs), 50):
        batch_docs = docs[j:j+50]
        batch_ids = ids[j:j+50]
        data = json.dumps({"content": batch_docs}).encode()
        
        for retry in range(3):
            try:
                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):
                    embs = [item["embedding"][0] for item in body]
                else:
                    embs = body.get("embedding", [])
                
                collection.add(ids=batch_ids, documents=batch_docs, embeddings=embs)
                print(".", end="", flush=True)
                break
            except Exception as e:
                if retry < 2:
                    print(f"R{retry}", end="", flush=True)
                    time.sleep(10)
                else:
                    print(f"X({e})", end="", flush=True)
        
        time.sleep(0.3)
    
    print(f" total={collection.count()}")

print(f"\nDone. Total: {collection.count()} documents")