#!/usr/bin/env python3
"""Create a self-contained OPD-only model-only recovery recipe.

DeepSpeed ZeRO-2 optimizer partitions are tied to the original world size, so
the four-rank checkpoint cannot be passed to ``Trainer.train(resume...)`` on
another world size.  The student model itself is portable.  This utility validates a
base recipe, replaces only the student source and output directory, and writes
the fully resolved YAML consumed by the new run.
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from typing import Any

import yaml


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from config.loader import load_config


def _require_complete_student_checkpoint(path: Path) -> Path:
    checkpoint = path.expanduser().resolve(strict=False)
    if not checkpoint.is_dir():
        raise FileNotFoundError(f"Student initialization checkpoint is not a directory: {checkpoint}")
    if not (checkpoint / "config.json").is_file():
        raise FileNotFoundError(f"Student initialization checkpoint lacks config.json: {checkpoint}")
    weight_files = tuple(checkpoint.glob("*.safetensors")) + tuple(checkpoint.glob("*.bin"))
    if not weight_files:
        raise FileNotFoundError(
            f"Student initialization checkpoint lacks model weights (*.safetensors or *.bin): {checkpoint}"
        )
    return checkpoint


def _materialize(
    *,
    base_config: Path,
    student_checkpoint: Path,
    output_dir: Path,
    output_config: Path,
    num_gpus: int,
    gradient_accumulation_steps: int | None,
) -> dict[str, Any]:
    """Return and atomically write the resolved model-only recovery recipe."""
    config = load_config(str(base_config))
    checkpoint = _require_complete_student_checkpoint(student_checkpoint)
    run_output = output_dir.expanduser().resolve(strict=False)
    target = output_config.expanduser().resolve(strict=False)

    if checkpoint == run_output or checkpoint in run_output.parents:
        raise ValueError(
            "The new output directory must not contain the model-initialization checkpoint; "
            "use a separate output directory for model-only recovery."
        )

    training = config["training"]
    dyme_args = training["dyme_args"]
    model = config["model"]

    # These are deliberately the only training-recipe changes.  All OPD loss,
    # filtering, evaluation, and checkpoint-evaluation settings stay exactly
    # those of the validated four-rank recipe.
    model["pretrained_model_path"] = str(checkpoint)
    model["teacher_device_map"] = "local"
    training["num_gpus"] = num_gpus
    training["num_client"] = num_gpus
    dyme_args["output_dir"] = str(run_output)
    if gradient_accumulation_steps is not None:
        dyme_args["gradient_accumulation_steps"] = gradient_accumulation_steps
    # A model-only process must not attempt to restore the four-rank ZeRO-2
    # optimizer/RNG partitions even if a future base config supplies one.
    dyme_args.pop("resume_from_checkpoint", None)

    target.parent.mkdir(parents=True, exist_ok=True)
    temporary = target.with_name(f".{target.name}.tmp-{os.getpid()}")
    with temporary.open("w", encoding="utf-8") as handle:
        yaml.safe_dump(config, handle, allow_unicode=True, sort_keys=False)
    os.replace(temporary, target)
    return config


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Materialize a self-contained OPD-only model-only recovery YAML."
    )
    parser.add_argument("--base-config", type=Path, required=True)
    parser.add_argument("--student-checkpoint", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--output-config", type=Path, required=True)
    parser.add_argument("--num-gpus", type=int, default=3)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=None)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.num_gpus <= 0:
        raise ValueError(f"--num-gpus must be positive, got {args.num_gpus}")
    if args.gradient_accumulation_steps is not None and args.gradient_accumulation_steps <= 0:
        raise ValueError(
            "--gradient-accumulation-steps must be positive when provided, "
            f"got {args.gradient_accumulation_steps}"
        )
    config = _materialize(
        base_config=args.base_config,
        student_checkpoint=args.student_checkpoint,
        output_dir=args.output_dir,
        output_config=args.output_config,
        num_gpus=args.num_gpus,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
    )
    print(
        f"Materialized {args.num_gpus}-rank model-only OPD recovery config: "
        f"student={config['model']['pretrained_model_path']} "
        f"output_dir={config['training']['dyme_args']['output_dir']} "
        f"config={args.output_config.expanduser().resolve(strict=False)}",
        flush=True,
    )


if __name__ == "__main__":
    main()