from __future__ import annotations

import argparse
import shutil
import tempfile
from pathlib import Path

from timeout_risk_model import load_dataset_records, save_model, train_timeout_model

DEFAULT_DATASET_REPO_ID = "stack-all/FJNU-HuaxiangChaPandaTimeoutRecord"


def default_dataset_dir() -> Path:
    return Path(__file__).resolve().parent.parent / "FJNU-HuaxiangChaPandaTimeoutRecord"


def download_dataset_snapshot(repo_id: str) -> Path:
    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise RuntimeError("Install huggingface_hub or use --dataset-dir with a local dataset checkout.") from exc

    target_dir = Path(tempfile.mkdtemp(prefix="cha_panda_dataset_"))
    snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        local_dir=str(target_dir),
        allow_patterns=["records.csv", "timeout_records.csv", "tmp/*.csv"],
    )
    return target_dir


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train the Huaxiang ChaPanda timeout risk model.")
    parser.add_argument(
        "--dataset-dir",
        type=Path,
        default=default_dataset_dir(),
        help="Local checkout of the Dataset repository.",
    )
    parser.add_argument(
        "--dataset-repo-id",
        default=DEFAULT_DATASET_REPO_ID,
        help="Hugging Face Dataset repo id used with --download-from-hub.",
    )
    parser.add_argument(
        "--download-from-hub",
        action="store_true",
        help="Download timeout_records.csv and tmp/*.csv from the Dataset repository before training.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path(__file__).resolve().parent / "timeout_model.json",
        help="Output JSON model path.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    downloaded_dir: Path | None = None
    dataset_dir = args.dataset_dir

    if args.download_from_hub:
        downloaded_dir = download_dataset_snapshot(args.dataset_repo_id)
        dataset_dir = downloaded_dir

    try:
        records = load_dataset_records(dataset_dir)
        model = train_timeout_model(records)
        save_model(model, args.output)
    finally:
        if downloaded_dir:
            shutil.rmtree(downloaded_dir, ignore_errors=True)

    metadata = model["metadata"]
    print(f"Saved model: {args.output}")
    print(f"Mode: {model['mode']}")
    print(f"Training records: {metadata['training_records']}")
    print(f"Timeout / non-timeout: {metadata['timeout_records']} / {metadata['non_timeout_records']}")
    print(f"Calibration: {metadata['calibration']}")


if __name__ == "__main__":
    main()