src/buckets.py

"""HF Storage Bucket helpers: list available CC crawls and validate the target.

All functions degrade gracefully: listing falls back to a static crawl list, and
validation returns a (ok, message) tuple instead of raising, so the UI can show a
friendly message.
"""
from __future__ import annotations

import re
from typing import List, Optional, Tuple

from . import config

_CRAWL_RE = re.compile(r"(CC-MAIN-\d{4}-\d{2})")


def _item_path(item) -> str:
    for attr in ("path", "name", "key"):
        val = getattr(item, attr, None)
        if isinstance(val, str) and val:
            return val
    return str(item)


def _crawls_under(api, prefix: str, token: Optional[str]) -> set:
    """Crawl ids found among the immediate children of `prefix` on the CC bucket."""
    crawls = set()
    for item in api.list_bucket_tree(config.CC_BUCKET, prefix=prefix, recursive=False, token=token):
        m = _CRAWL_RE.search(_item_path(item))
        if m:
            crawls.add(m.group(1))
    return crawls


def list_available_crawls(token: Optional[str] = None) -> List[str]:
    """Crawl ids usable for repackaging, newest first.

    A crawl is only usable if BOTH its WARC data (``crawl-data/``) and its columnar
    index (``cc-index/table/cc-main/warc/``) are present, so we return the
    intersection of crawls found under both prefixes. Falls back to
    ``config.FALLBACK_CRAWLS`` on any error (offline, auth, API change)."""
    try:
        from huggingface_hub import HfApi

        api = HfApi()
        data_crawls = _crawls_under(api, "crawl-data", token)
        index_crawls = _crawls_under(api, config.CC_INDEX_SUBPATH.rstrip("/"), token)
        usable = data_crawls & index_crawls
        if usable:
            return sorted(usable, reverse=True)
    except Exception:
        pass
    return list(config.FALLBACK_CRAWLS)


# Org roles that grant write access to the org's buckets.
_WRITE_ROLES = {"admin", "write", "contributor"}


def list_writable_buckets(token: Optional[str] = None) -> List[str]:
    """Bucket ids ("namespace/bucket") the signed-in user can write to.

    Covers the user's own buckets plus buckets in organizations where they have a
    write role. Returns [] on any error (offline, not signed in)."""
    if not token:
        return []
    try:
        from huggingface_hub import HfApi

        api = HfApi()
        me = api.whoami(token=token)
        namespaces = [me["name"]]
        namespaces += [
            o["name"] for o in me.get("orgs", []) if o.get("roleInOrg") in _WRITE_ROLES
        ]
        ids = set()
        for ns in namespaces:
            try:
                for b in api.list_buckets(namespace=ns, token=token):
                    bid = getattr(b, "id", None)
                    if bid:
                        ids.add(bid)
            except Exception:  # noqa: BLE001 - skip a namespace we can't list
                continue
        return sorted(ids)
    except Exception:  # noqa: BLE001
        return []


def parse_targets(text: str) -> List[str]:
    """Split a free-text list of hostnames/domains (commas, whitespace, newlines)."""
    if not text:
        return []
    parts = re.split(r"[\s,]+", text.strip())
    return [p.strip().lower() for p in parts if p.strip()]


# Hostnames/domains only contain these characters (also what cdx_toolkit accepts).
_HOST_RE = re.compile(r"^[A-Za-z0-9.\-]+$")


def validate_targets(hostnames: List[str], domains: List[str]) -> Tuple[bool, str]:
    if not hostnames and not domains:
        return False, "Enter at least one hostname or domain to filter for."
    for value in hostnames + domains:
        if not _HOST_RE.match(value) or "." not in value:
            return False, f"Invalid hostname/domain: {value!r} (letters, digits, dot, hyphen only)."
    return True, ""


def validate_languages(languages: List[str]) -> Tuple[bool, str]:
    """Language codes must come from the hardcoded ISO-639-3 list (empty = any)."""
    unknown = [c for c in (languages or []) if c not in config.CONTENT_LANGUAGE_CODES]
    if unknown:
        return False, "Unknown ISO-639-3 language code(s): " + ", ".join(unknown)
    return True, ""


def validate_crawls(crawls: List[str]) -> Tuple[bool, str]:
    if not crawls:
        return False, "Select at least one crawl."
    if len(crawls) > config.COST_GUARD_MAX_CRAWLS:
        return (
            True,
            f"⚠️ {len(crawls)} crawls selected (> {config.COST_GUARD_MAX_CRAWLS}). The index scan "
            "may be large and slow; consider fewer crawls.",
        )
    return True, ""


def validate_target_bucket(
    bucket_id: str, path: str, token: Optional[str] = None
) -> Tuple[bool, str]:
    """Confirm the target is an existing HF **bucket** (not a dataset/model repo).

    Dataset repos are not supported as a target (HF Jobs cannot mount/write a
    dataset the way it mounts a bucket), which `bucket_info` enforces naturally:
    it only resolves bucket ids.
    """
    bucket_id = (bucket_id or "").strip().strip("/")
    if bucket_id.count("/") != 1:
        return False, "Target bucket must be of the form `namespace/bucket`."
    if any(c in (path or "") for c in ("..",)) or (path or "").startswith("/"):
        return False, "Target path must be a relative sub-path inside the bucket."
    try:
        from huggingface_hub import HfApi

        HfApi().bucket_info(bucket_id, token=token)
    except Exception as e:  # noqa: BLE001 - surface a friendly message
        return (
            False,
            f"Could not find a writable bucket `{bucket_id}`. It must exist before launching "
            f"(dataset repos are not supported as a target). Details: {type(e).__name__}: {e}",
        )
    return True, f"Target bucket `{bucket_id}` found."