#!/usr/bin/env python3
"""Ingest JSONL into ChromaDB — uses llama.cpp /embedding endpoint."""
import json, os, sys, glob, urllib.request, time
import chromadb

LLAMA_URL = "http://localhost:8081/embedding"

def llama_embed(texts):
    """Get embeddings from nomic-embed server with retry."""
    data = json.dumps({"content": texts}).encode()
    req = urllib.request.Request(LLAMA_URL, data=data,
        headers={"Content-Type": "application/json"})
    for attempt in range(3):
        try:
            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:
            if attempt < 2:
                print(f"[retry {attempt+1}] {e}")
                time.sleep(5)
            else:
                raise


class LlamaEmbedFn:
    """Custom embedding function using llama.cpp."""
    def __call__(self, input):
        all_embs = []
        for i in range(0, len(input), 32):
            batch = input[i:i+32]
            embs = llama_embed(batch)
            if isinstance(embs[0], float):
                all_embs.append(embs)
            else:
                all_embs.extend(embs)
        return all_embs


def main():
    jsonl_dir = sys.argv[1] if len(sys.argv) > 1 else os.path.expanduser("~/nubes/data/elsa_jsonl_wi_en")
    db_dir = sys.argv[2] if len(sys.argv) > 2 else os.path.expanduser("~/nubes/chroma_db")
    collection_name = sys.argv[3] if len(sys.argv) > 3 else "elsa_docs"
    
    ef = LlamaEmbedFn()
    
    client = chromadb.PersistentClient(path=db_dir)
    
    try:
        client.delete_collection(collection_name)
    except:
        pass
    
    collection = client.create_collection(
        name=collection_name,
        embedding_function=ef,
        metadata={"hnsw:space": "cosine"}
    )
    
    jsonl_files = sorted(glob.glob(os.path.join(jsonl_dir, "*.jsonl")))
    print(f"Found {len(jsonl_files)} JSONL files")
    print(f"Using llama.cpp at {LLAMA_URL}")
    
    total = 0
    for jf in jsonl_files:
        ids, docs, metas = [], [], []
        
        with open(jf, "r", encoding="utf-8") as f:
            for line in f:
                try:
                    d = json.loads(line)
                    title = d.get("title", "").strip()
                    text = d.get("text", "").strip()
                    ft = d.get("full_title", "").strip()
                    
                    if ft and text:
                        doc = f"{ft}\n{text}"
                    elif title and text:
                        doc = f"{title}\n{text}"
                    else:
                        doc = text or title or ft
                    
                    if not doc or len(doc) < 20:
                        continue
                    
                    # Split long docs into chunks of ~1000 chars
                    chunks = [doc[i:i+1000] for i in range(0, len(doc), 1000)]
                    for chunk in chunks:
                        if len(chunk) < 20:
                            continue
                        ids.append(str(total))
                        docs.append(chunk)
                        metas.append({
                            "source": d.get("source", ""),
                            "file": d.get("file", ""),
                            "title": title or ft or doc[:100],
                        })
                        total += 1
                except:
                    continue
        
        if ids:
            # Send in smaller batches of 500
            for i in range(0, len(ids), 500):
                batch_ids = ids[i:i+500]
                batch_docs = docs[i:i+500]
                batch_metas = metas[i:i+500]
                for attempt in range(3):
                    try:
                        collection.add(ids=batch_ids, documents=batch_docs, metadatas=batch_metas)
                        break
                    except Exception as e:
                        if attempt < 2:
                            print(f"  [retry {os.path.basename(jf)} batch {attempt+1}] {e}")
                            time.sleep(10)
                        else:
                            print(f"  [SKIP batch] {e}")
            print(f"  {os.path.basename(jf)}: +{len(ids)} (total {total})")
    
    print(f"\nDone. {total} documents → {db_dir}/{collection_name}")

if __name__ == "__main__":
    main()