"""
Search HuggingFace Hub for pre-curated Indian heritage art datasets.

HuggingFace hosts user-contributed datasets that are perfect for LoRA training.
This script searches for relevant datasets and lets you download them in one shot.

Usage:
    python scripts/search_hf_datasets.py
    python scripts/search_hf_datasets.py --download madhubani-art
"""
from __future__ import annotations

import argparse
import logging
from pathlib import Path

log = logging.getLogger(__name__)

# Known datasets worth checking (will verify and update at runtime)
KNOWN_DATASETS = [
    # Indian art / painting datasets
    "rishiMadhuri/madhubani-art",
    "MMadhubani/madhubani-painting",
    "ravi105/madhubani-art",
    "nerdlab/indian-art",
    "shubhamorhan/indian-painting-styles",
    "Sakshi1307/indian-art-forms",
    "DrishtiSharma/indian-heritage-art",
    "rcss/indian-art-classification",
    # General art datasets that include Indian art
    "huggan/few-shot-art",
    "huggan/wikiart",
    "Artificio/WikiArt",
    # Bharat-specific
    "Bharat/desi-art",
    "kabita-chatterjee/indian-art-forms",
]

# Search queries to run on the HF Hub API
SEARCH_QUERIES = [
    "madhubani",
    "warli",
    "pattachitra",
    "mughal painting",
    "tanjore painting",
    "indian art",
    "indian heritage",
    "indian painting",
    "folk art india",
    "traditional art india",
]


def search_hf_datasets(query: str, limit: int = 10) -> list:
    """Search HuggingFace Hub for datasets matching a query."""
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        results = list(api.list_datasets(search=query, limit=limit))
        return [
            {
                "id": d.id,
                "downloads": d.downloads,
                "likes": d.likes,
                "tags": d.tags,
            }
            for d in results
        ]
    except Exception as exc:
        log.error(f"HF search failed for '{query}': {exc}")
        return []


def check_dataset_contents(dataset_id: str) -> dict:
    """Check if a dataset has images and what splits/columns it has."""
    try:
        from datasets import load_dataset_builder
        builder = load_dataset_builder(dataset_id)
        info = {
            "id": dataset_id,
            "description": (builder.info.description or "")[:200],
            "features": str(builder.info.features)[:500],
            "splits": list(builder.info.splits.keys()) if builder.info.splits else [],
            "size": builder.info.dataset_size or "unknown",
        }
        return info
    except Exception as exc:
        return {"id": dataset_id, "error": str(exc)[:200]}


def download_dataset(dataset_id: str, out_dir: Path, max_images: int = 50) -> int:
    """Download a dataset and save images to out_dir. Returns count saved."""
    out_dir.mkdir(parents=True, exist_ok=True)
    try:
        from datasets import load_dataset
        ds = load_dataset(dataset_id, split="train")
        saved = 0
        for i, item in enumerate(ds):
            if saved >= max_images:
                break
            # Find the image field
            img = None
            for key in ("image", "img", "painting", "art", "picture"):
                if key in item:
                    img = item[key]
                    break
            if img is None:
                # Try first PIL Image field
                for k, v in item.items():
                    if hasattr(v, "save"):
                        img = v
                        break
            if img is None:
                continue
            try:
                img = img.convert("RGB")
                out_path = out_dir / f"{dataset_id.replace('/', '_')}_{saved:03d}.jpg"
                img.save(out_path, "JPEG", quality=95)
                saved += 1
            except Exception:
                continue
        return saved
    except Exception as exc:
        log.error(f"Download failed for {dataset_id}: {exc}")
        return 0


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")

    p = argparse.ArgumentParser()
    p.add_argument("--query", help="Search for datasets matching this query")
    p.add_argument("--download", help="Download this dataset ID")
    p.add_argument("--style", help="Save downloads to assets/datasets/raw/<style>/")
    p.add_argument("--max", type=int, default=50, help="Max images to download")
    p.add_argument("--check", help="Check contents of this dataset ID")
    args = p.parse_args()

    if args.check:
        info = check_dataset_contents(args.check)
        print(f"\n=== {args.check} ===")
        for k, v in info.items():
            print(f"  {k}: {v}")
        return

    if args.download:
        out_dir = Path(f"assets/datasets/raw/{args.style}") if args.style else Path("outputs/hf_download")
        n = download_dataset(args.download, out_dir, max_images=args.max)
        print(f"Downloaded {n} images to {out_dir}")
        return

    # Default: search across all queries
    print("=" * 60)
    print("HuggingFace Heritage Art Dataset Search")
    print("=" * 60)
    all_results = {}

    for q in SEARCH_QUERIES:
        print(f"\n--- Search: '{q}' ---")
        results = search_hf_datasets(q, limit=5)
        for r in results:
            print(f"  {r['id']:50s} | ↓{r['downloads']:>6} | ♥{r['likes']:>3} | tags: {[t for t in r.get('tags', []) if 'task' in t][:2]}")
            all_results[r["id"]] = r

    # Also check known datasets
    print("\n--- Checking known datasets ---")
    for ds_id in KNOWN_DATASETS:
        info = check_dataset_contents(ds_id)
        if "error" in info:
            print(f"  ✗ {ds_id}: {info['error'][:80]}")
        else:
            print(f"  ✓ {ds_id}")
            print(f"      splits: {info['splits']}")
            print(f"      features: {info['features'][:200]}")

    print(f"\n=== Found {len(all_results)} candidate datasets ===")
    print("To check a dataset's contents:")
    print(f"  python scripts/search_hf_datasets.py --check <dataset-id>")
    print("\nTo download a dataset:")
    print(f"  python scripts/search_hf_datasets.py --download <dataset-id> --style madhubani --max 50")


if __name__ == "__main__":
    main()