#!/usr/bin/env python3
"""Inspect a Hugging Face dataset using the Dataset Viewer API.

This mirrors the useful parts of upstream ml-intern's hf_inspect_dataset tool:
status, configs/splits, schema, sample rows, and parquet file availability.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any


BASE_URL = "https://datasets-server.huggingface.co"
MAX_SAMPLE_VALUE_LEN = 150


def fetch_json(path: str, params: dict[str, Any], token: str | None) -> dict[str, Any]:
    query = urllib.parse.urlencode({k: v for k, v in params.items() if v is not None})
    request = urllib.request.Request(f"{BASE_URL}{path}?{query}")
    if token:
        request.add_header("Authorization", f"Bearer {token}")
    try:
        with urllib.request.urlopen(request, timeout=30) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"{path} returned HTTP {exc.code}: {body[:500]}") from exc


def type_name(feature: Any) -> str:
    if isinstance(feature, str):
        return feature
    if not isinstance(feature, dict):
        return type(feature).__name__
    feature_type = feature.get("_type")
    if feature_type == "ClassLabel":
        names = feature.get("names") or []
        if 0 < len(names) <= 5:
            values = ", ".join(f"{name}={idx}" for idx, name in enumerate(names))
            return f"ClassLabel ({values})"
        return f"ClassLabel ({len(names)} classes)"
    if feature_type:
        return feature_type
    if "dtype" in feature:
        return str(feature["dtype"])
    return json.dumps(feature, ensure_ascii=False)[:120]


def extract_configs(splits_data: dict[str, Any]) -> list[dict[str, Any]]:
    configs: dict[str, dict[str, Any]] = {}
    for item in splits_data.get("splits", []):
        config = item.get("config", "default")
        split = item.get("split", "train")
        row_count = item.get("num_rows") or item.get("num_examples")
        configs.setdefault(config, {"name": config, "splits": []})
        configs[config]["splits"].append({"name": split, "rows": row_count})
    return list(configs.values())


def format_status(data: dict[str, Any]) -> str:
    available = [
        key
        for key in ("viewer", "preview", "search", "filter", "statistics")
        if data.get(key)
    ]
    if available:
        return f"## Status\nValid ({', '.join(available)})"
    return "## Status\nDataset may have Dataset Viewer issues"


def format_structure(configs: list[dict[str, Any]], max_rows: int = 20) -> str:
    lines = ["## Structure (configs & splits)", "| Config | Split | Rows |", "|---|---|---:|"]
    total = sum(len(config["splits"]) for config in configs)
    shown = 0
    for config in configs:
        for split in config["splits"]:
            if shown >= max_rows:
                break
            rows = split["rows"] if split["rows"] is not None else "?"
            lines.append(f"| {config['name']} | {split['name']} | {rows} |")
            shown += 1
        if shown >= max_rows:
            break
    if total > shown:
        lines.append("| ... | ... | ... |")
        lines.append(f"_Showing {shown} of {total} config/split rows._")
    return "\n".join(lines)


def format_schema(info: dict[str, Any], config: str) -> str:
    features = info.get("dataset_info", {}).get("features", {})
    lines = [f"## Schema ({config})", "| Column | Type |", "|---|---|"]
    if not features:
        lines.append("| (none found) | unknown |")
    for column, feature in features.items():
        lines.append(f"| {column} | {type_name(feature)} |")
    return "\n".join(lines)


def maybe_json(value: Any) -> Any:
    if isinstance(value, str):
        try:
            return json.loads(value)
        except json.JSONDecodeError:
            return value
    return value


def format_messages(messages: Any) -> str | None:
    messages = maybe_json(messages)
    if not isinstance(messages, list) or not messages:
        return None
    roles: set[str] = set()
    keys: set[str] = set()
    has_tool_calls = False
    has_tool_results = False
    example: dict[str, Any] | None = None
    fallback: dict[str, Any] | None = None
    for message in messages:
        if not isinstance(message, dict):
            continue
        keys.update(message.keys())
        if message.get("role"):
            roles.add(str(message["role"]))
        if message.get("tool_calls") or message.get("function_call"):
            has_tool_calls = True
            example = example or message
        if message.get("role") in {"tool", "function"} or message.get("tool_call_id"):
            has_tool_results = True
        if message.get("role") == "assistant":
            example = example or message
        elif message.get("role") != "system":
            fallback = fallback or message
    example = example or fallback
    lines = ["## Messages Column Format"]
    lines.append(f"Roles: {', '.join(sorted(roles)) if roles else 'unknown'}")
    common = ["role", "content", "tool_calls", "tool_call_id", "name", "function_call"]
    lines.append("Message keys: " + ", ".join(f"{key} {'yes' if key in keys else 'no'}" for key in common))
    if has_tool_calls:
        lines.append("Tool calls: present")
    if has_tool_results:
        lines.append("Tool results: present")
    if example:
        cleaned = dict(example)
        content = cleaned.get("content")
        if isinstance(content, str) and len(content) > 100:
            cleaned["content"] = content[:100] + "..."
        lines.append("")
        lines.append("Example message structure:")
        lines.append("```json")
        lines.append(json.dumps(cleaned, indent=2, ensure_ascii=False))
        lines.append("```")
    return "\n".join(lines)


def format_samples(rows_data: dict[str, Any], config: str, split: str, limit: int) -> str:
    rows = rows_data.get("rows", [])[:limit]
    lines = [f"## Sample Rows ({config}/{split})"]
    first_messages: Any = None
    for idx, row_wrapper in enumerate(rows, 1):
        row = row_wrapper.get("row", {})
        lines.append(f"**Row {idx}:**")
        for key, value in row.items():
            if key.lower() == "messages" and first_messages is None:
                first_messages = value
            text = str(value)
            if len(text) > MAX_SAMPLE_VALUE_LEN:
                text = text[:MAX_SAMPLE_VALUE_LEN] + "..."
            lines.append(f"- {key}: {text}")
    if not rows:
        lines.append("(no rows returned)")
    message_section = format_messages(first_messages) if first_messages is not None else None
    if message_section:
        lines.append("")
        lines.append(message_section)
    return "\n".join(lines)


def format_parquet(data: dict[str, Any], max_rows: int = 20) -> str | None:
    files = data.get("parquet_files", [])
    if not files:
        return None
    groups: dict[str, dict[str, int]] = {}
    for item in files:
        key = f"{item.get('config', 'default')}/{item.get('split', 'train')}"
        groups.setdefault(key, {"count": 0, "size": 0})
        groups[key]["count"] += 1
        size = item.get("size") or 0
        groups[key]["size"] += int(size) if isinstance(size, (int, float)) else 0
    lines = ["## Files (Parquet)"]
    for key, values in list(groups.items())[:max_rows]:
        size_mb = values["size"] / (1024 * 1024)
        lines.append(f"- {key}: {values['count']} file(s), {size_mb:.1f} MB")
    if len(groups) > max_rows:
        lines.append(f"- ... showing {max_rows} of {len(groups)} groups")
    return "\n".join(lines)


def compatibility_notes(features: dict[str, Any]) -> str:
    columns = set(features)
    lines = ["## Training Compatibility"]
    checks = {
        "SFT": bool({"messages", "text"} & columns or {"prompt", "completion"} <= columns),
        "DPO": {"prompt", "chosen", "rejected"} <= columns,
        "GRPO": "prompt" in columns,
    }
    for method, ok in checks.items():
        lines.append(f"- {method}: {'looks compatible' if ok else 'columns not sufficient'}")
    if "messages" in columns:
        lines.append("- Chat data: inspect the sample message roles before choosing a trainer template.")
    return "\n".join(lines)


def inspect_dataset(dataset: str, config: str | None, split: str | None, sample_rows: int, token: str | None) -> str:
    warnings: list[str] = []
    with ThreadPoolExecutor(max_workers=3) as pool:
        futures = {
            pool.submit(fetch_json, "/is-valid", {"dataset": dataset}, token): "is-valid",
            pool.submit(fetch_json, "/splits", {"dataset": dataset}, token): "splits",
            pool.submit(fetch_json, "/parquet", {"dataset": dataset}, token): "parquet",
        }
        phase1: dict[str, Any] = {}
        for future in as_completed(futures):
            name = futures[future]
            try:
                phase1[name] = future.result()
            except Exception as exc:
                warnings.append(f"{name}: {exc}")

    configs = extract_configs(phase1.get("splits", {}))
    selected_config = config or (configs[0]["name"] if configs else "default")
    selected_split = split or (configs[0]["splits"][0]["name"] if configs and configs[0]["splits"] else "train")

    with ThreadPoolExecutor(max_workers=2) as pool:
        futures = {
            pool.submit(fetch_json, "/info", {"dataset": dataset, "config": selected_config}, token): "info",
            pool.submit(
                fetch_json,
                "/first-rows",
                {"dataset": dataset, "config": selected_config, "split": selected_split},
                token,
            ): "first-rows",
        }
        phase2: dict[str, Any] = {}
        for future in as_completed(futures):
            name = futures[future]
            try:
                phase2[name] = future.result()
            except Exception as exc:
                warnings.append(f"{name}: {exc}")

    features = phase2.get("info", {}).get("dataset_info", {}).get("features", {})
    sections = [f"# {dataset}"]
    if "is-valid" in phase1:
        sections.append(format_status(phase1["is-valid"]))
    if configs:
        sections.append(format_structure(configs))
    if "info" in phase2:
        sections.append(format_schema(phase2["info"], selected_config))
        sections.append(compatibility_notes(features))
    if "first-rows" in phase2:
        sections.append(format_samples(phase2["first-rows"], selected_config, selected_split, sample_rows))
    parquet = format_parquet(phase1.get("parquet", {}))
    if parquet:
        sections.append(parquet)
    if warnings:
        sections.append("## Warnings\n" + "\n".join(f"- {warning}" for warning in warnings))
    return "\n\n".join(sections)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("dataset", help="Dataset id, for example stanfordnlp/imdb")
    parser.add_argument("--config", help="Config/subset name")
    parser.add_argument("--split", help="Split for sample rows")
    parser.add_argument("--sample-rows", type=int, default=3, help="Number of rows to show, max 10")
    parser.add_argument("--token-env", default="HF_TOKEN", help="Environment variable containing an HF token")
    args = parser.parse_args()
    token = os.environ.get(args.token_env)
    print(inspect_dataset(args.dataset, args.config, args.split, min(args.sample_rows, 10), token))
    return 0


if __name__ == "__main__":
    sys.exit(main())