dataset_loader.py

"""
dataset_loader.py — TEKDEV Bot Personality Dataset Fetcher

SETUP (do this once):
  1. Go to https://huggingface.co/new-dataset
  2. Name it:  gemma3-personality
  3. Set visibility: Private (recommended)
  4. Upload any of these .md files to the repo root:
       identity.md   ← who the bot is
       tone.md       ← communication style
       rules.md      ← hard guardrails
       knowledge.md  ← domain hints  (optional)
       examples.md   ← few-shot samples (optional)

  Your dataset will live at:
    https://huggingface.co/datasets/YOUR_USERNAME/gemma3-personality

  5. Add  HF_USERNAME=your_username  as a Space secret.
     The full dataset ID resolves to:  your_username/gemma3-personality

This script downloads all .md files from that dataset into
PERSONALITY_CACHE_DIR so personality.py can read them.
"""

import logging
import shutil
from pathlib import Path

from huggingface_hub import hf_hub_download, list_repo_files
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError

from config import HF_TOKEN, PERSONALITY_DATASET, PERSONALITY_CACHE_DIR

logger = logging.getLogger(__name__)

# Direct browse URL (informational, printed in logs)
_DATASET_URL = (
    f"https://huggingface.co/datasets/{PERSONALITY_DATASET}"
    if PERSONALITY_DATASET
    else "https://huggingface.co/new-dataset"
)


def _ensure_cache_dir() -> Path:
    cache = Path(PERSONALITY_CACHE_DIR)
    cache.mkdir(parents=True, exist_ok=True)
    return cache


def fetch_personality_files(force_refresh: bool = False) -> list[str]:
    """
    Download all *.md files from:
      https://huggingface.co/datasets/{PERSONALITY_DATASET}

    Args:
        force_refresh: Re-download even if local copies exist.

    Returns:
        List of filenames successfully cached to PERSONALITY_CACHE_DIR.
    """
    if not PERSONALITY_DATASET:
        logger.error(
            "PERSONALITY_DATASET is empty — cannot fetch personality files.\n"
            "Set HF_USERNAME in your Space secrets so the dataset ID resolves to:\n"
            "  YOUR_USERNAME/gemma3-personality\n"
            "Create the dataset at: https://huggingface.co/new-dataset"
        )
        return []

    logger.info("Fetching personality files from: %s", _DATASET_URL)
    cache = _ensure_cache_dir()

    # List files in the dataset repo
    try:
        all_files = list(
            list_repo_files(
                repo_id=PERSONALITY_DATASET,
                repo_type="dataset",
                token=HF_TOKEN or None,
            )
        )
    except RepositoryNotFoundError:
        logger.error(
            "Dataset not found: %s\n"
            "  → Check HF_USERNAME is correct\n"
            "  → Confirm the dataset exists at %s\n"
            "  → If private, make sure HF_TOKEN has read access",
            PERSONALITY_DATASET, _DATASET_URL,
        )
        return []
    except Exception as exc:
        logger.error("Failed to list dataset '%s': %s", PERSONALITY_DATASET, exc)
        return []

    md_files = [f for f in all_files if f.endswith(".md")]

    if not md_files:
        logger.warning(
            "No .md files found in %s\n"
            "  → Upload identity.md, tone.md, rules.md etc. to the dataset root at:\n"
            "  → %s",
            PERSONALITY_DATASET, _DATASET_URL,
        )
        return []

    fetched: list[str] = []

    for filename in md_files:
        dest = cache / Path(filename).name
        if dest.exists() and not force_refresh:
            logger.debug("Cached already: %s", filename)
            fetched.append(dest.name)
            continue

        try:
            tmp_path = hf_hub_download(
                repo_id=PERSONALITY_DATASET,
                filename=filename,
                repo_type="dataset",
                token=HF_TOKEN or None,
                force_download=force_refresh,
            )
            shutil.copy2(tmp_path, dest)
            logger.info("  ✓ %s", filename)
            fetched.append(dest.name)
        except EntryNotFoundError:
            logger.warning("  ✗ Not found in dataset: %s", filename)
        except Exception as exc:
            logger.error("  ✗ Error downloading %s: %s", filename, exc)

    logger.info(
        "Personality sync done — %d/%d files | dataset: %s",
        len(fetched), len(md_files), _DATASET_URL,
    )
    return fetched


def clear_cache() -> None:
    cache = Path(PERSONALITY_CACHE_DIR)
    if cache.exists():
        shutil.rmtree(cache)
        logger.info("Cache cleared: %s", cache)


def refresh() -> list[str]:
    """Wipe cache and re-download everything from the dataset."""
    clear_cache()
    return fetch_personality_files(force_refresh=True)


if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.INFO)
    files = fetch_personality_files(force_refresh="--refresh" in sys.argv)
    print(f"\nCached: {files}")
    print(f"Dataset URL: {_DATASET_URL}")