#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import random
import time
from typing import Iterable
import urllib.error
import urllib.parse
import urllib.request

from huggingface_hub import HfApi


DEFAULT_DATASETS = {
    "english_web_fineweb": {
        "repo_id": "HuggingFaceFW/fineweb",
        "patterns": (".parquet",),
        "default_max_files": 256,
    },
    "english_edu_fineweb_edu": {
        "repo_id": "HuggingFaceFW/fineweb-edu",
        "patterns": (".parquet",),
        "default_max_files": 256,
    },
    "chinese_cci3_hq": {
        "repo_id": "BAAI/CCI3-HQ",
        "patterns": (".jsonl", ".parquet"),
        "default_max_files": 128,
    },
    "chinese_skypile_150b": {
        "repo_id": "Skywork/SkyPile-150B",
        "patterns": (".jsonl", ".parquet"),
        "default_max_files": 128,
    },
    "science_medmcqa": {
        "repo_id": "openlifescienceai/medmcqa",
        "patterns": (".parquet",),
        "default_max_files": None,
    },
    "science_proofwriter": {
        "repo_id": "tasksource/proofwriter",
        "patterns": (".parquet",),
        "default_max_files": None,
    },
    "science_scienceqa": {
        "repo_id": "derek-thomas/ScienceQA",
        "patterns": (".parquet",),
        "default_max_files": None,
    },
    "science_sciq": {
        "repo_id": "allenai/sciq",
        "patterns": (".parquet",),
        "default_max_files": None,
    },
    "science_qasc": {
        "repo_id": "allenai/qasc",
        "patterns": (".parquet",),
        "default_max_files": None,
    },
    "science_openbookqa": {
        "repo_id": "allenai/openbookqa",
        "patterns": (".parquet",),
        "default_max_files": None,
    },
    "ms_fineweb_edu": {
        "provider": "modelscope",
        "repo_id": "AI-ModelScope/fineweb-edu",
        "roots": ("data",),
        "patterns": (".parquet",),
        "default_max_files": 96,
    },
    "ms_ultra_fineweb_en": {
        "provider": "modelscope",
        "repo_id": "OpenBMB/Ultra-FineWeb",
        "roots": ("data/ultrafineweb_en",),
        "patterns": (".parquet",),
        "default_max_files": 128,
    },
    "ms_ultra_fineweb_zh": {
        "provider": "modelscope",
        "repo_id": "OpenBMB/Ultra-FineWeb",
        "roots": ("data/ultrafineweb_zh",),
        "patterns": (".parquet",),
        "default_max_files": 32,
    },
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Download rebalanced pretraining source datasets via hf-mirror.")
    parser.add_argument("--output-root", default="/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--endpoint", default=os.environ.get("HF_ENDPOINT", "https://hf-mirror.com"))
    parser.add_argument("--modelscope-endpoint", default=os.environ.get("MODELSCOPE_ENDPOINT", "https://modelscope.cn"))
    parser.add_argument("--token", default=os.environ.get("HF_TOKEN"))
    parser.add_argument("--include", action="append", default=[], help="Dataset key to include. Repeatable.")
    parser.add_argument("--exclude", action="append", default=[], help="Dataset key to skip. Repeatable.")
    parser.add_argument("--max-files", action="append", default=[], metavar="KEY=N")
    parser.add_argument("--manifest-name", default="download_manifest.jsonl")
    parser.add_argument("--socket-timeout", type=float, default=60.0)
    parser.add_argument(
        "--max-mbps",
        type=float,
        default=0.0,
        help="Per-process download bandwidth cap in MiB/s. 0 disables throttling.",
    )
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.token:
        raise ValueError("HF_TOKEN is required")
    output_root = Path(args.output_root)
    output_root.mkdir(parents=True, exist_ok=True)
    max_files = parse_max_files(args.max_files)
    selected_keys = set(args.include) if args.include else set(DEFAULT_DATASETS)
    selected_keys -= set(args.exclude)
    unknown = selected_keys - set(DEFAULT_DATASETS)
    if unknown:
        raise ValueError(f"unknown dataset keys: {sorted(unknown)}")

    api = HfApi(endpoint=args.endpoint, token=args.token)
    manifest_path = output_root / args.manifest_name
    completed = load_completed_downloads(manifest_path)
    for key in sorted(selected_keys):
        spec = DEFAULT_DATASETS[key]
        repo_id = spec["repo_id"]
        provider = spec.get("provider", "hf")
        limit = max_files.get(key, spec["default_max_files"])
        if provider == "modelscope":
            files = list_modelscope_dataset_files(
                endpoint=args.modelscope_endpoint,
                repo_id=repo_id,
                roots=spec.get("roots", ("",)),
                suffixes=spec["patterns"],
                max_files=limit,
                seed=args.seed + stable_int(key),
            )
        else:
            files = list_dataset_files(api, repo_id=repo_id, suffixes=spec["patterns"])
        random.Random(args.seed + stable_int(key)).shuffle(files)
        if limit is not None:
            files = files[:limit]
        dataset_dir = output_root / key
        dataset_dir.mkdir(parents=True, exist_ok=True)
        active_endpoint = args.modelscope_endpoint if provider == "modelscope" else args.endpoint
        print(
            "download_dataset_start "
            f"key={key} repo_id={repo_id} provider={provider} files={len(files)} "
            f"output={dataset_dir} endpoint={active_endpoint}",
            flush=True,
        )
        for index, filename in enumerate(files, start=1):
            completed_path = completed.get((key, filename))
            if completed_path is not None and completed_path.exists():
                record = {
                    "time": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
                    "key": key,
                    "repo_id": repo_id,
                    "filename": filename,
                    "index": index,
                    "total": len(files),
                    "status": "skipped_existing",
                    "elapsed_sec": 0.0,
                    "local_path": str(completed_path),
                    "seed": args.seed,
                    "error": "",
                }
                append_jsonl(manifest_path, record)
                print(
                    "download_file "
                    f"key={key} index={index}/{len(files)} status=skipped_existing file={filename}",
                    flush=True,
                )
                continue
            started_at = time.time()
            status = "ok"
            error = ""
            local_path = ""
            if args.dry_run:
                status = "dry_run"
                local_path = str(dataset_dir / filename)
            else:
                try:
                    local_path = download_file_with_resume(
                        repo_id=repo_id,
                        filename=filename,
                        output_dir=dataset_dir,
                        token=args.token,
                        endpoint=args.modelscope_endpoint if provider == "modelscope" else args.endpoint,
                        provider=provider,
                        timeout=args.socket_timeout,
                        max_mbps=args.max_mbps,
                    )
                except Exception as exc:  # keep overnight batch moving
                    status = "error"
                    error = f"{type(exc).__name__}: {exc}"
            elapsed = time.time() - started_at
            record = {
                "time": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
                "key": key,
                "repo_id": repo_id,
                "provider": provider,
                "filename": filename,
                "index": index,
                "total": len(files),
                "status": status,
                "elapsed_sec": elapsed,
                "local_path": local_path,
                "seed": args.seed,
                "error": error,
            }
            append_jsonl(manifest_path, record)
            print(
                "download_file "
                f"key={key} index={index}/{len(files)} status={status} file={filename} elapsed={elapsed:.1f}s",
                flush=True,
            )
        print(f"download_dataset_done key={key} files={len(files)}", flush=True)


def list_dataset_files(api: HfApi, *, repo_id: str, suffixes: Iterable[str]) -> list[str]:
    suffixes = tuple(suffixes)
    info = api.dataset_info(repo_id)
    files = []
    for sibling in info.siblings or []:
        name = sibling.rfilename
        if name.endswith(suffixes):
            files.append(name)
    if not files:
        raise FileNotFoundError(f"no files matched {suffixes} in {repo_id}")
    return sorted(files)


def download_file_with_resume(
    *,
    repo_id: str,
    filename: str,
    output_dir: Path,
    token: str,
    endpoint: str,
    provider: str,
    timeout: float,
    max_mbps: float,
) -> str:
    target = output_dir / filename
    target.parent.mkdir(parents=True, exist_ok=True)
    url = resolve_url(endpoint=endpoint, repo_id=repo_id, filename=filename, provider=provider)
    remote_size = remote_content_length(url, token=token, timeout=timeout)
    if target.exists() and remote_size is not None and target.stat().st_size == remote_size:
        return str(target)
    incomplete = target.with_name(target.name + ".incomplete")
    offset = incomplete.stat().st_size if incomplete.exists() else 0
    if remote_size is not None and offset > remote_size:
        incomplete.unlink()
        offset = 0
    headers = {"Authorization": f"Bearer {token}"}
    if offset:
        headers["Range"] = f"bytes={offset}-"
    request = urllib.request.Request(url, headers=headers)
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            mode = "ab" if offset and response.status == 206 else "wb"
            throttle = DownloadThrottle(max_mbps)
            with incomplete.open(mode) as handle:
                while True:
                    chunk = response.read(1024 * 1024)
                    if not chunk:
                        break
                    handle.write(chunk)
                    throttle.wait(len(chunk))
    except urllib.error.HTTPError as exc:
        if exc.code == 416 and remote_size is not None and incomplete.exists() and incomplete.stat().st_size == remote_size:
            incomplete.replace(target)
            return str(target)
        raise
    if remote_size is not None and incomplete.stat().st_size != remote_size:
        raise IOError(f"incomplete download for {filename}: got {incomplete.stat().st_size}, expected {remote_size}")
    incomplete.replace(target)
    return str(target)


class DownloadThrottle:
    def __init__(self, max_mbps: float) -> None:
        self.max_bytes_per_sec = max(0.0, max_mbps) * 1024 * 1024
        self.started_at = time.monotonic()
        self.bytes_seen = 0

    def wait(self, bytes_read: int) -> None:
        if self.max_bytes_per_sec <= 0:
            return
        self.bytes_seen += bytes_read
        expected_elapsed = self.bytes_seen / self.max_bytes_per_sec
        actual_elapsed = time.monotonic() - self.started_at
        sleep_sec = expected_elapsed - actual_elapsed
        if sleep_sec > 0:
            time.sleep(sleep_sec)


def resolve_url(*, endpoint: str, repo_id: str, filename: str, provider: str) -> str:
    base = endpoint.rstrip("/")
    quoted_filename = urllib.parse.quote(filename, safe="/")
    if provider == "modelscope":
        return f"{base}/datasets/{repo_id}/resolve/master/{quoted_filename}"
    return f"{base}/datasets/{repo_id}/resolve/main/{quoted_filename}"


def remote_content_length(url: str, *, token: str, timeout: float) -> int | None:
    request = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"}, method="HEAD")
    with urllib.request.urlopen(request, timeout=timeout) as response:
        raw = response.headers.get("Content-Length")
    return int(raw) if raw and raw.isdigit() else None


def list_modelscope_dataset_files(
    *,
    endpoint: str,
    repo_id: str,
    roots: Iterable[str],
    suffixes: Iterable[str],
    max_files: int | None,
    seed: int,
) -> list[str]:
    files: list[str] = []
    suffixes = tuple(suffixes)
    for root in roots:
        remaining = None if max_files is None else max_files - len(files)
        if remaining is not None and remaining <= 0:
            break
        files.extend(
            _list_modelscope_tree(
                endpoint=endpoint,
                repo_id=repo_id,
                root=root,
                suffixes=suffixes,
                max_files=remaining,
                seed=seed + stable_int(root),
            )
        )
    if not files:
        raise FileNotFoundError(f"no ModelScope files matched {suffixes} in {repo_id} roots={tuple(roots)}")
    return sorted(set(files))


def _list_modelscope_tree(
    *,
    endpoint: str,
    repo_id: str,
    root: str,
    suffixes: tuple[str, ...],
    max_files: int | None,
    seed: int,
) -> list[str]:
    files: list[str] = []
    page = 1
    page_size = 500
    trees: list[str] = []
    while True:
        query = urllib.parse.urlencode(
            {
                "Revision": "master",
                "Root": root,
                "PageNumber": page,
                "PageSize": page_size,
            }
        )
        url = f"{endpoint.rstrip('/')}/api/v1/datasets/{repo_id}/repo/tree?{query}"
        request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0", "Accept": "application/json"})
        with urllib.request.urlopen(request, timeout=60) as response:
            payload = json.load(response)
        data = payload.get("Data")
        entries = data.get("Files") if isinstance(data, dict) else data
        if not entries:
            break
        for entry in entries:
            path = entry.get("Path")
            entry_type = entry.get("Type")
            if entry_type == "blob" and isinstance(path, str) and path.endswith(suffixes):
                files.append(path)
                if max_files is not None and len(files) >= max_files:
                    return files
            elif entry_type == "tree" and isinstance(path, str):
                trees.append(path)
        total = int(payload.get("TotalCount") or (data or {}).get("TotalCount") or 0)
        if page * page_size >= total:
            break
        page += 1
    random.Random(seed).shuffle(trees)
    for tree_root in trees:
        remaining = None if max_files is None else max_files - len(files)
        if remaining is not None and remaining <= 0:
            break
        files.extend(
            _list_modelscope_tree(
                endpoint=endpoint,
                repo_id=repo_id,
                root=tree_root,
                suffixes=suffixes,
                max_files=remaining,
                seed=seed + stable_int(tree_root),
            )
        )
    return files


def parse_max_files(items: list[str]) -> dict[str, int | None]:
    parsed: dict[str, int | None] = {}
    for item in items:
        if "=" not in item:
            raise ValueError(f"--max-files must use KEY=N: {item!r}")
        key, raw = item.split("=", 1)
        key = key.strip()
        raw = raw.strip().lower()
        if raw in {"all", "none", "-1"}:
            parsed[key] = None
        else:
            value = int(raw)
            if value <= 0:
                raise ValueError(f"max files must be positive: {item!r}")
            parsed[key] = value
    return parsed


def append_jsonl(path: Path, record: dict[str, object]) -> None:
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")


def load_completed_downloads(path: Path) -> dict[tuple[str, str], Path]:
    completed: dict[tuple[str, str], Path] = {}
    if not path.exists():
        return completed
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            try:
                record = json.loads(line)
            except json.JSONDecodeError:
                continue
            if record.get("status") != "ok":
                continue
            key = record.get("key")
            filename = record.get("filename")
            local_path = record.get("local_path")
            if isinstance(key, str) and isinstance(filename, str) and isinstance(local_path, str):
                completed[(key, filename)] = Path(local_path)
    return completed


def stable_int(value: str) -> int:
    import hashlib

    digest = hashlib.sha256(value.encode("utf-8")).digest()
    return int.from_bytes(digest[:4], "big")


if __name__ == "__main__":
    main()
