#!/usr/bin/env python3
# coding: utf-8
"""
Quick diagnosis for Ministral startup "hang" caused by invalid model/tokenizer path.

Usage:
  python chengyue_script/ministral/diagnose_ministral_model_path.py
  python chengyue_script/ministral/diagnose_ministral_model_path.py \
      --model-path /path/to/model --tokenizer-path /path/to/model
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path
from typing import Any, Callable


def _fmt_exc(exc: BaseException) -> str:
    msg = str(exc).strip()
    return msg if msg else repr(exc)


def _time_call(name: str, fn: Callable[[], Any]) -> None:
    t0 = time.time()
    print(f"\n[CHECK] {name}")
    try:
        out = fn()
        dt = time.time() - t0
        print(f"[OK] {name} succeeded in {dt:.2f}s: {type(out).__name__}")
    except Exception as exc:  # noqa: BLE001
        dt = time.time() - t0
        print(f"[FAIL] {name} failed in {dt:.2f}s")
        print(f"       {type(exc).__name__}: {_fmt_exc(exc)}")


def main() -> int:
    parser = argparse.ArgumentParser(description="Diagnose Ministral model/tokenizer path issues.")
    parser.add_argument(
        "--model-path",
        default="src/models/Nemotron-Diffusion-Exp-Ministral-8B-Instruct-merged",
        help="Model path passed to --model_name_or_path in training script.",
    )
    parser.add_argument(
        "--tokenizer-path",
        default="src/models/Nemotron-Diffusion-Exp-Ministral-8B-Instruct-merged",
        help="Tokenizer path passed to --tokenizer_name in training script.",
    )
    parser.add_argument(
        "--trust-remote-code",
        type=int,
        default=1,
        help="Whether to pass trust_remote_code.",
    )
    args = parser.parse_args()

    print("=== Environment ===")
    print(f"python: {sys.executable}")
    print(f"cwd: {os.getcwd()}")
    print(f"model_path(raw): {args.model_path}")
    print(f"tokenizer_path(raw): {args.tokenizer_path}")

    model_path = Path(args.model_path)
    tokenizer_path = Path(args.tokenizer_path)
    print(f"model_path(abs): {model_path.resolve()}")
    print(f"tokenizer_path(abs): {tokenizer_path.resolve()}")
    print(f"model_path.exists(): {model_path.exists()}")
    print(f"tokenizer_path.exists(): {tokenizer_path.exists()}")
    if not model_path.exists() or not tokenizer_path.exists():
        print(
            "\n[WARN] Local path does not exist. In multi-rank deepspeed startup, "
            "this often looks like a long hang before all ranks finally error out."
        )

    from transformers import AutoConfig, AutoTokenizer

    trust_remote_code = bool(args.trust_remote_code)

    # 1) Strict local check: should fail fast if local path is invalid.
    _time_call(
        "AutoTokenizer(local_files_only=True)",
        lambda: AutoTokenizer.from_pretrained(
            args.tokenizer_path,
            trust_remote_code=trust_remote_code,
            local_files_only=True,
        ),
    )
    _time_call(
        "AutoConfig(local_files_only=True)",
        lambda: AutoConfig.from_pretrained(
            args.model_path,
            trust_remote_code=trust_remote_code,
            local_files_only=True,
        ),
    )

    # 2) Default behavior: may attempt hub/cache resolution and be slower.
    _time_call(
        "AutoTokenizer(default local_files_only=False)",
        lambda: AutoTokenizer.from_pretrained(
            args.tokenizer_path,
            trust_remote_code=trust_remote_code,
        ),
    )
    _time_call(
        "AutoConfig(default local_files_only=False)",
        lambda: AutoConfig.from_pretrained(
            args.model_path,
            trust_remote_code=trust_remote_code,
        ),
    )

    print("\n=== Diagnosis Hint ===")
    print(
        "If local_files_only=True already fails, fix --model_name_or_path/--tokenizer_name first.\n"
        "A valid local model dir must contain at least config/tokenizer files."
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

