from __future__ import annotations

"""
Lightweight planning helpers for the `hfschema wizard` CLI.

These functions are intentionally pure and do not import `warpdata` or touch
the network so they can be tested easily. The idea is:

- Given a warp dataset name (e.g. "dmt-brains" or "vision/coco"), suggest
  what kind of HF dataset(s) should exist and how they relate.
- The interactive CLI (implemented in cli.py) can call these helpers and
  then render human-readable instructions or commands.
"""

from dataclasses import dataclass
from enum import Enum
from typing import List


class PlanKind(str, Enum):
    """High-level migration pattern for a dataset."""

    PLAIN_TABLE = "plain_table"
    TABLE_AND_RAW = "table_and_raw"
    TABLE_AND_SIGNALS = "table_and_signals"
    IMAGE_WITH_IMAGES = "image_with_images"
    METADATA_ONLY = "metadata_only"
    URL_METADATA_ONLY = "url_metadata_only"


@dataclass
class WarpDatasetPlan:
    """Recommended HF-facing plan for a given warp dataset."""

    warp_name: str
    pattern: PlanKind
    hf_ids: List[str]
    notes: str = ""
    # Optional: warp dataset name to use when resolving local Parquet paths.
    # For most datasets this matches `warp_name`, but for cases like
    # vision/imagenet-1k we may want to use vision/imagenet-1k-images as the
    # source for image paths.
    source_warp_name: str | None = None


def _normalize_warp_name(name: str) -> str:
    """
    Normalize a warp dataset name to a simple slug.

    We drop any leading "warpdata://" and workspace prefixes so that
    "warpdata://vision/coco" and "coco" both map to "coco".
    """
    # Drop URI prefix if present.
    if name.startswith("warpdata://"):
        name = name[len("warpdata://") :]
    # Keep only the final segment after workspace/ prefixes.
    if "/" in name:
        name = name.split("/")[-1]
    return name.strip()


def suggest_plan_for_warp_dataset(name: str) -> WarpDatasetPlan:
    """
    Suggest a migration pattern + HF ids for a warp dataset.

    This function is driven by hard-coded knowledge of common "hard"
    datasets in the author's warp environment. It is deliberately simple
    and deterministic so it can be tested without depending on live
    caches.
    """
    slug = _normalize_warp_name(name)

    # EEG: DMT Brains – metadata + numeric signals.
    if slug in {
        "dmt-brains",
        "dmt-brains-baseline",
        "dmt-brains-complete",
        "dmt-brains-conditions",
        "dmt-brains-dmt-only",
    }:
        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.TABLE_AND_SIGNALS,
            hf_ids=["uwunion/dmt-brains-conditions", "uwunion/dmt-brains-signals"],
            notes=(
                "EEG dataset: keep subject/condition metadata and windowed signals "
                "in separate HF datasets, linked by subject_id + condition."
            ),
        )

    # Audio: VCTK – utterances with Audio() + speakers metadata.
    if slug in {"vctk", "vctk-speakers"}:
        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.TABLE_AND_SIGNALS,
            hf_ids=["uwunion/vctk-utterances", "uwunion/vctk-speakers"],
            notes=(
                "Speech corpus: utterance-level dataset with Audio() plus a "
                "separate speakers metadata table."
            ),
        )

    # Mocap: CMU – metadata + raw BVH repo.
    if slug in {"cmu", "cmu-subjects"}:
        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.TABLE_AND_RAW,
            hf_ids=["uwunion/cmu-mocap", "uwunion/cmu-mocap-raw"],
            notes=(
                "MoCap dataset: publish motion-level metadata and keep BVH files "
                "in a dedicated raw-file HF repo referenced via hf_path."
            ),
        )

    # Vision: COCO / CelebA / ImageNet images – upload images too.
    if slug in {
        "celeba-attrs",
        "coco",
        "coco-train",
        "coco-val",
        "coco-test",
        "imagenet-1k",
        "imagenet-1k-images",
    }:
        hf_id = {
            "celeba-attrs": "uwunion/celeba-attrs",
            "coco": "uwunion/coco-mixed",
            "coco-train": "uwunion/coco-train",
            "coco-val": "uwunion/coco-val",
            "coco-test": "uwunion/coco-test",
            # For ImageNet, always prefer the images-backed HF dataset.
            "imagenet-1k": "uwunion/imagenet-1k-images",
            "imagenet-1k-images": "uwunion/imagenet-1k-images",
        }[slug]

        # For imagenet-1k, the warp dataset that actually has image_path is
        # usually vision/imagenet-1k-images. If the user passed a dataset
        # name containing "imagenet-1k", prefer a "-images" variant when
        # resolving local Parquet paths.
        source_warp_name: str | None = None
        if slug == "imagenet-1k":
            if "imagenet-1k" in name:
                source_warp_name = name.replace("imagenet-1k", "imagenet-1k-images")
            else:
                source_warp_name = "imagenet-1k-images"

        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.IMAGE_WITH_IMAGES,
            hf_ids=[hf_id],
            notes=(
                "Vision dataset with local image paths: publish as an HF dataset "
                "with an Image() column so images are stored on HF."
            ),
            source_warp_name=source_warp_name,
        )

    # Vision: SAM 3D Body – metadata only, rely on canonical HF dataset.
    if slug == "sam-3d-body":
        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.METADATA_ONLY,
            hf_ids=["uwunion/sam-3d-body-metadata"],
            notes=(
                "SAM 3D Body is already fully hosted as facebook/sam-3d-body-dataset; "
                "keep your merged 3D metadata as a separate HF table and refer back "
                "to the canonical dataset via IDs."
            ),
        )

    # Vision: LAION – URL-based images; treat the table (including URLs) as
    # the full dataset. We do not attempt to upload all underlying images.
    if slug in {"laion-10m"}:
        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.PLAIN_TABLE,
            hf_ids=["uwunion/laion-10m"],
            notes=(
                "LAION subset uses remote image URLs; publish the full table "
                "to HF (including URLs) so any machine can load text+URL and "
                "fetch images on demand."
            ),
        )

    # Finance: stocks/options/earnings/news/fundamentals/crypto – plain tables.
    if any(
        slug.startswith(prefix)
        for prefix in (
            "stocks-",
            "hourly-ohlc",
            "news-",
            "sec-",
            "fundamentals-",
            "sectors",
            "options-",
            "spy-options-",
        )
    ):
        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.PLAIN_TABLE,
            hf_ids=[f"uwunion/{slug}"],
            notes="Financial dataset: publish as a plain Parquet-backed HF table.",
        )

    if any(slug.startswith(prefix) for prefix in ("binance/", "coingecko/")):
        # Strip workspace prefix if present.
        simple = slug.split("/")[-1]
        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.PLAIN_TABLE,
            hf_ids=[f"uwunion/{simple}"],
            notes="Crypto/market dataset: publish as a plain Parquet-backed HF table.",
        )

    # Text / math / logic that are "hard" only because of size, not assets.
    if slug in {
        "mathx-100k",
        "numina_lean",
        "wikipedia-main",
        "fineweb2-mix",
        "logician",
        "logician-v2",
        "mathvista",
        "mathvision",
        "math_hendrycks",
        "mmlu",
        "truthfulqa",
        "pubmedqa-artificial",
        "pubmedqa-labeled",
    }:
        return WarpDatasetPlan(
            warp_name=name,
            pattern=PlanKind.PLAIN_TABLE,
            hf_ids=[f"uwunion/{slug.replace('_', '-')}"],
            notes="Large text/math/logic corpus: publish as a plain HF table.",
        )

    # Default: treat as a plain table mapped to uwunion/<slug>.
    return WarpDatasetPlan(
        warp_name=name,
        pattern=PlanKind.PLAIN_TABLE,
        hf_ids=[f"uwunion/{slug}"],
        notes="Default plan: publish as a single Parquet-backed HF table.",
    )
