#!/usr/bin/env python3
# scripts/ingest_public_data.py
# ============================================
# RUN ONCE: python scripts/ingest_public_data.py
# Populates ChromaDB with defect patch triplets
# ============================================

import chromadb
from sentence_transformers import SentenceTransformer
from PIL import Image
import json
from pathlib import Path
import argparse


# ============================================
# DATASET REGISTRY — Add new sources here
# ============================================

DATASETS = {
    # ── Commercial-friendly datasets ──
    "roboflow_manufacturing": {
        "path": "data/external/roboflow_manufacturing",
        "domain": "pcb",
        "license": "cc-by-4.0",
        "commercial_ok": True,
        "attribution": "Roboflow Universe — Manufacturing Defect Detection",
        "description": "PCB manufacturing defects, CC BY 4.0 licensed"
    },
    
    # ── Your proprietary data ──
    "my_vcsel_captures": {
        "path": "data/my_captures/vcsel",
        "domain": "glass/vcsel",
        "license": "proprietary",
        "commercial_ok": True,
        "attribution": "Internal",
        "description": "VCSEL laser diode defect captures"
    },
    
    # ── Partner data (NDA) ──
    # "partner_pcb_nda": {
    #     "path": "data/partners/pcb_client_a",
    #     "domain": "pcb",
    #     "license": "partner_nda",
    #     "commercial_ok": True,
    #     "attribution": "Partner NDA",
    #     "description": "PCB client defect images under NDA"
    # },
}


def validate_triplet(folder: Path) -> dict | None:
    """Check if folder contains required patch triplet files."""
    required = ["original_masked.png", "original_target.png", "artifact_target.png"]
    paths = {f: folder / f for f in required}
    
    for f, p in paths.items():
        if not p.exists():
            print(f"  ⚠️  Missing {f} in {folder}, skipping")
            return None
    
    return {k: str(v) for k, v in paths.items()}


def generate_caption(folder_name: str, domain: str, dataset_desc: str) -> str:
    """Generate a text caption for embedding."""
    # You can replace this with a VLM call for better captions
    defect_type = folder_name.replace("_", " ")
    return f"{defect_type} defect on {domain}. {dataset_desc}"


def ingest_dataset(collection, encoder, name: str, cfg: dict):
    """Ingest one dataset into ChromaDB."""
    print(f"\n{'='*50}")
    print(f"Dataset: {name}")
    print(f"License: {cfg['license']} | Commercial: {'✅' if cfg['commercial_ok'] else '❌'}")
    print(f"Path: {cfg['path']}")
    print(f"{'='*50}")
    
    base_path = Path(cfg["path"])
    if not base_path.exists():
        print(f"  ⚠️  Path not found: {base_path}")
        print(f"  Create it and add patch triplet folders:")
        print(f"    {base_path}/defect_name_001/original_masked.png")
        print(f"    {base_path}/defect_name_001/original_target.png")
        print(f"    {base_path}/defect_name_001/artifact_target.png")
        return 0

    # NEW: Get set of IDs already in the collection to avoid re-processing
    existing_ids = set(collection.get()["ids"])

    count = 0
    for triplet_folder in sorted(base_path.iterdir()):
        if not triplet_folder.is_dir():
            continue

        # NEW: Skip if already processed
        doc_id = f"{name}_{triplet_folder.name}"
        if doc_id in existing_ids:
            print(f"  ⏭️  Skipping {triplet_folder.name} (already in DB)")
            continue
        
        paths = validate_triplet(triplet_folder)
        if paths is None:
            continue
        
        caption = generate_caption(
            triplet_folder.name,
            cfg["domain"],
            cfg.get("description", "")
        )
        
        embedding = encoder.encode(caption)
        
        collection.upsert(
            documents=[caption],
            embeddings=[embedding.tolist()],
            metadatas=[{
                "paths": json.dumps(paths),
                "domain": cfg["domain"],
                "license": cfg["license"],
                "commercial_ok": cfg["commercial_ok"],
                "source": name,
                "attribution": cfg["attribution"],
                "defect_name": triplet_folder.name
            }],
            ids=[f"{name}_{triplet_folder.name}"]
        )
        
        count += 1
        print(f" ✅  {triplet_folder.name}: {caption[:60]}...")
    
    return count


def main():
    parser = argparse.ArgumentParser(description="Ingest defect patch triplets into RAG DB")
    parser.add_argument("--db-path", default="data/defect_db", help="ChromaDB persistent path")
    parser.add_argument("--collection", default="defect_patches", help="Collection name")
    parser.add_argument("--model", default="all-MiniLM-L6-v2", help="SentenceTransformer model")
    args = parser.parse_args()
    
    # Initialize DB
    Path(args.db_path).mkdir(parents=True, exist_ok=True)
    client = chromadb.PersistentClient(path=args.db_path)
    
    # Delete existing collection if you want fresh start
    # client.delete_collection(args.collection)
    
    collection = client.get_or_create_collection(
        name=args.collection,
        metadata={"hnsw:space": "cosine"}
    )
    
    print(f"DB path: {args.db_path}")
    print(f"Collection: {args.collection}")
    print(f"Existing entries: {collection.count()}")
    
    # Initialize encoder
    print(f"\nLoading encoder: {args.model}")
    encoder = SentenceTransformer(args.model)
    
    # Ingest all datasets
    total = 0
    for name, cfg in DATASETS.items():
        # Skip non-commercial datasets in commercial builds
        if not cfg.get("commercial_ok", False):
            print(f"\n⏭️  Skipping {name} — not commercial-friendly")
            continue
        
        count = ingest_dataset(collection, encoder, name, cfg)
        total += count
    
    print(f"\n{'='*50}")
    print(f"TOTAL INGESTED: {total} patch triplets")
    print(f"TOTAL IN DB: {collection.count()}")
    print(f"{'='*50}")
    
    # Print commercial summary
    print("\n📋 Commercial License Summary:")
    results = collection.get()
    licenses = {}
    for meta in results["metadatas"]:
        lic = meta["license"]
        licenses[lic] = licenses.get(lic, 0) + 1
    
    for lic, count in licenses.items():
        icon = "✅" if any(
            cfg["license"] == lic and cfg.get("commercial_ok")
            for cfg in DATASETS.values()
        ) else "❌"
        print(f"  {icon} {lic}: {count} entries")


if __name__ == "__main__":
    main()