#!/usr/bin/env python3
"""Train from one LLaMA-Factory YAML on four GPUs; import Unsloth in each worker first."""
from __future__ import annotations

import argparse
from datetime import datetime, timezone
import functools
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
from types import MethodType

ROOT = Path(__file__).resolve().parents[1]
GPUS = 4


def read_config(path: Path) -> dict:
    import yaml
    with path.open(encoding="utf-8-sig") as handle:
        config = yaml.safe_load(handle)
    if not isinstance(config, dict):
        raise ValueError("The configuration must be a YAML mapping.")
    if config.get("use_unsloth") is not True:
        raise ValueError("This package requires use_unsloth: true.")
    if config.get("bf16") is not True or config.get("fp16", False):
        raise ValueError("The H200 configuration must use bf16, not fp16.")
    if "model_max_length" in config:
        raise ValueError("Remove model_max_length from YAML; LLaMA-Factory derives it from cutoff_len.")
    if int(config.get("cutoff_len", 0)) < 1:
        raise ValueError("cutoff_len must be positive.")
    for key in ("dataset_dir", "output_dir"):
        if key not in config:
            raise ValueError(f"Missing configuration key: {key}")
        value = Path(config[key]).expanduser()
        config[key] = str((ROOT / value).resolve() if not value.is_absolute() else value.resolve())
    return config


def write_config(path: Path, config: dict) -> None:
    import yaml
    path.write_text(yaml.safe_dump(config, allow_unicode=True, sort_keys=False), encoding="utf-8")


def digest(path: Path) -> str:
    result = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            result.update(block)
    return result.hexdigest()


def verify_data(data_dir: Path, splits: tuple[str, ...]) -> dict:
    manifest_path = data_dir / "manifest.json"
    if not manifest_path.is_file():
        raise ValueError(f"Run prepare_data.py first; missing {manifest_path}")
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    for split in splits:
        if manifest.get("splits", {}).get(split, {}).get("rows", 0) < 1:
            raise ValueError(f"The prepared {split} split is empty.")
    names = ["dataset_info.json"]
    for split in splits:
        names.extend((f"{split}.jsonl", f"{split}.metadata.jsonl"))
    for name in names:
        item = manifest.get("files", {}).get(name, {})
        if not item.get("sha256") or digest(data_dir / name) != item["sha256"]:
            raise ValueError(f"Prepared data changed or is missing: {name}")
    return manifest


def run_logged(command: list[str], log: Path, env: dict | None = None) -> None:
    print("Running:", " ".join(command), flush=True)
    with log.open("a", encoding="utf-8") as handle:
        process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                   text=True, encoding="utf-8", errors="replace", env=env)
        try:
            for line in process.stdout:
                print(line, end="", flush=True)
                handle.write(line)
                handle.flush()
            code = process.wait()
        except BaseException:
            process.terminate()
            process.wait()
            raise
    if code:
        raise subprocess.CalledProcessError(code, command)


def context_settings(settings: dict, original_config) -> dict:
    requested = int(settings["cutoff_len"])
    native = int(original_config.max_position_embeddings)
    if requested > native or settings.get("rope_scaling"):
        raise ValueError(f"Use cutoff_len <= {native} within the model's native context and omit rope_scaling.")
    return {"requested_length": requested, "base_positions": native}


def reserve_generation_space() -> None:
    # Reserve space before batching so the trainer sees the same prompt length
    # as model.generate when it masks the prompt in the generated sequence.
    from llamafactory.data.processor import unsupervised

    original = unsupervised.infer_seqlen

    @functools.wraps(original)
    def with_generation_space(source_len: int, target_len: int, cutoff_len: int):
        source, target = original(source_len, target_len, cutoff_len)
        source = min(source, cutoff_len - 1)
        if source < source_len or target < target_len:
            print(f"Generation preprocessing truncated tokens: input {source_len} -> {source}, "
                  f"label {target_len} -> {target}, context {cutoff_len}.", flush=True)
        return source, target

    unsupervised.infer_seqlen = with_generation_space


def bound_generation_context(model, context_length: int, default_max_new_tokens: int) -> None:
    original_generate = model.generate

    @functools.wraps(original_generate)
    def bounded_generate(self, *args, **kwargs):
        inputs = kwargs.get("input_ids")
        if inputs is None:
            inputs = kwargs.get("inputs")
        if inputs is None and args:
            inputs = args[0]
        if inputs is None:
            raise ValueError("Cannot verify the text input length before generation.")
        prompt_length = int(inputs.shape[-1])
        available = context_length - prompt_length
        if available < 1:
            raise ValueError("The preprocessed input leaves no room for an answer inside the context limit.")
        requested = kwargs.get("max_new_tokens")
        requested = int(default_max_new_tokens if requested is None else requested)
        if requested < 1:
            raise ValueError("max_new_tokens must be positive.")
        kwargs["max_new_tokens"] = min(requested, available)
        if requested > available:
            print(f"Generation limit: {available} new tokens after {prompt_length} input tokens "
                  f"(context {context_length}).", flush=True)
        return original_generate(*args, **kwargs)

    model.generate = MethodType(bounded_generate, model)


def require_unsloth_backend(settings: dict, context: dict) -> dict:
    # LLaMA-Factory may otherwise silently fall back when Unsloth rejects a model.
    # The wrappers affect this process only; no installed package is edited.
    from llamafactory.model import loader
    from llamafactory.model.model_utils import unsloth as integration
    loaded = {}

    if hasattr(integration, "_get_unsloth_kwargs"):
        original_kwargs = integration._get_unsloth_kwargs

        @functools.wraps(original_kwargs)
        def pinned_kwargs(config, model_name_or_path, model_args, finetuning_args):
            values = original_kwargs(config, model_name_or_path, model_args, finetuning_args)
            values["use_exact_model_name"] = True
            if model_args.model_revision:
                values["revision"] = model_args.model_revision
            return values

        integration._get_unsloth_kwargs = pinned_kwargs

    def strict(original):
        @functools.wraps(original)
        def wrapped(*args, **kwargs):
            model = original(*args, **kwargs)
            if model is None:
                raise RuntimeError("Unsloth did not load this model. Training/evaluation stopped; "
                                   "use a compatible Unsloth/LLaMA-Factory environment.")
            if model.config.max_position_embeddings < context["requested_length"]:
                raise RuntimeError("The loaded model's configured context is shorter than the requested SFT context.")
            if settings.get("do_train"):
                loaded["base_config_json"] = model.config.to_json_string()
            return model
        return wrapped

    for name in ("load_unsloth_pretrained_model", "load_unsloth_peft_model"):
        original = getattr(loader, name, None)
        if original is None:
            raise RuntimeError("The installed LLaMA-Factory has an unsupported Unsloth integration.")
        setattr(loader, name, strict(original))
    if settings.get("do_predict") and settings.get("predict_with_generate"):
        # Keep PEFT/Unsloth delegation: GenerationMixin.generate bound directly
        # to a PeftModel fails in Transformers 5 because its class has no _sample.
        original_patch_model = loader.patch_model

        @functools.wraps(original_patch_model)
        def patch_model_preserving_generate(model, *args, **kwargs):
            original_generate = model.generate
            result = original_patch_model(model, *args, **kwargs)
            model.generate = original_generate
            return result

        loader.patch_model = patch_model_preserving_generate
        # Attach the guard to the final model, after LLaMA-Factory patches generate
        # and finishes adapter initialization. Earlier attachment can be replaced.
        from llamafactory.train.sft import workflow
        original_load_model = workflow.load_model

        @functools.wraps(original_load_model)
        def load_bounded_model(*args, **kwargs):
            model = original_load_model(*args, **kwargs)
            bound_generation_context(model, context["requested_length"], int(settings["max_new_tokens"]))
            return model

        workflow.load_model = load_bounded_model
    return loaded


def worker(config_path: Path) -> None:
    # This order is important for Unsloth's patches.
    import unsloth  # noqa: F401
    import torch
    config = read_config(config_path)
    world = int(os.environ.get("WORLD_SIZE", "1"))
    if config.get("do_train") and world != GPUS:
        raise RuntimeError("This training scenario requires four processes / four GPUs.")
    if not torch.cuda.is_available() or torch.cuda.device_count() < world:
        raise RuntimeError(f"This job requires {world} visible CUDA GPU(s).")
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    torch.cuda.set_device(local_rank)
    if not torch.cuda.is_bf16_supported():
        raise RuntimeError("The allocated GPU does not support the requested bf16 precision.")
    print(f"rank={local_rank}, world_size={world}, GPU={torch.cuda.get_device_name(local_rank)}, "
          "precision=bf16, backend=Unsloth", flush=True)
    from transformers import AutoConfig
    original_config = AutoConfig.from_pretrained(config["model_name_or_path"], revision=config["model_revision"], trust_remote_code=False)
    context = context_settings(config, original_config)
    print("Model context:", json.dumps(context), flush=True)
    loaded = require_unsloth_backend(config, context)
    if config.get("do_predict") and config.get("predict_with_generate"):
        if context["requested_length"] < 2:
            raise ValueError("Generation requires cutoff_len >= 2.")
        reserve_generation_space()
    from llamafactory.train.tuner import run_exp
    run_exp(config)
    if config.get("do_train") and local_rank == 0:
        (Path(config["output_dir"]) / "base_model_config.json").write_text(loaded["base_config_json"], encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("config", type=Path, help="LLaMA-Factory model YAML.")
    parser.add_argument("--run-dir", type=Path, help="New run directory; adapter is saved inside adapter/.")
    parser.add_argument("--resume", type=Path, help="Checkpoint inside the original run's adapter/ directory.")
    parser.add_argument("--pilot-steps", type=int, help="Limit a pilot to this many optimizer steps; requires --run-dir.")
    parser.add_argument("--pilot-samples", type=int,
                        help="Maximum rows per dataset in a pilot (default: 2048 or all eval rows, whichever is larger).")
    parser.add_argument("--pilot-eval-steps", type=int,
                        help="Evaluate and save at this interval in a pilot (default: at the last step).")
    parser.add_argument("--dry-run", action="store_true", help="Validate data/config; do not import GPU libraries.")
    parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS)
    args = parser.parse_args()
    if args.worker:
        worker(args.config.resolve())
        return 0
    if args.pilot_steps is None and args.pilot_samples is not None:
        parser.error("--pilot-samples requires --pilot-steps.")
    if args.pilot_steps is not None and (args.pilot_steps < 1 or args.run_dir is None):
        parser.error("--pilot-steps must be positive and requires --run-dir.")
    if args.pilot_samples is not None and args.pilot_samples < 1:
        parser.error("--pilot-samples must be positive.")
    if args.pilot_eval_steps is not None and not 1 <= args.pilot_eval_steps <= (args.pilot_steps or 0):
        parser.error("--pilot-eval-steps requires --pilot-steps and must be between 1 and that limit.")
    config = read_config(args.config.resolve())
    if config.get("stage") != "sft" or config.get("finetuning_type") != "lora":
        raise ValueError("Expected an SFT / LoRA training configuration.")
    if not config.get("do_train") or config.get("eval_dataset") != "aya_eval":
        raise ValueError("Training must use the prepared train and eval splits.")
    if config.get("dataset") != "aya_train" or "val_size" in config:
        raise ValueError("Use aya_train and aya_eval without another automatic split.")
    if any(key in config for key in ("max_steps", "max_samples")):
        raise ValueError("The full-data scenario must not impose a smoke-test sample/step limit.")
    batch = GPUS * config["per_device_train_batch_size"] * config["gradient_accumulation_steps"]
    data_dir = Path(config["dataset_dir"])
    manifest = verify_data(data_dir, ("train", "eval"))
    train_rows = manifest["splits"]["train"]["rows"]
    eval_rows = manifest["splits"]["eval"]["rows"]
    if args.pilot_steps is not None:
        # LLaMA-Factory limits both splits before tokenization. Keep the complete
        # validation split, even if the caller prepared more than 2048 eval rows.
        sample_limit = args.pilot_samples or max(2048, eval_rows)
        if sample_limit < eval_rows:
            parser.error(f"--pilot-samples must be at least {eval_rows} to retain the full eval split.")
        eval_interval = args.pilot_eval_steps or args.pilot_steps
        config.update(max_steps=args.pilot_steps, max_samples=sample_limit,
                      eval_strategy="steps", save_strategy="steps",
                      eval_steps=eval_interval, save_steps=eval_interval)
    if args.run_dir:
        config["output_dir"] = str(args.run_dir.resolve() / "adapter")
    run_dir = Path(config["output_dir"]).parent
    if args.resume:
        checkpoint = args.resume.resolve()
        checkpoint.relative_to(Path(config["output_dir"]))
        for name in ("trainer_state.json", "adapter_config.json", "optimizer.pt", "scheduler.pt"):
            if not (checkpoint / name).is_file():
                raise ValueError(f"Incomplete checkpoint: {checkpoint / name}")
        previous = json.loads((run_dir / "run.json").read_text(encoding="utf-8"))
        if previous["config"] != config or previous["data_manifest_sha256"] != digest(data_dir / "manifest.json"):
            raise ValueError("Resume requires the original configuration and prepared data.")
    elif run_dir.exists() and any(run_dir.iterdir()):
        raise ValueError(f"Run directory is not empty; choose --run-dir for a new run: {run_dir}")
    summary = {"gpus_per_training_run": GPUS, "effective_batch": batch,
               "cutoff_len": config["cutoff_len"], "bf16": True, "use_unsloth": True,
               "train_rows": train_rows, "eval_rows": eval_rows,
               "adapter_output": config["output_dir"]}
    if args.pilot_steps is not None:
        summary.update(pilot_steps=args.pilot_steps, pilot_train_rows=min(train_rows, config["max_samples"]),
                       pilot_eval_steps=config["eval_steps"])
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    if args.dry_run:
        print("Configuration/data checks passed. No training started.")
        return 0
    run_dir.mkdir(parents=True, exist_ok=True)
    if not args.resume:
        record = {"config": config, "gpus": GPUS, "effective_batch": batch,
                  "data_manifest_sha256": digest(data_dir / "manifest.json"),
                  "created_at_utc": datetime.now(timezone.utc).isoformat()}
        (run_dir / "run.json").write_text(json.dumps(record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    execution = dict(config)
    if args.resume:
        execution["resume_from_checkpoint"] = str(args.resume.resolve())
    execution_path = run_dir / "used_config.yaml"
    write_config(execution_path, execution)
    env = os.environ.copy()
    env.update({"PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false"})
    env.setdefault("OMP_NUM_THREADS", "4")
    command = [sys.executable, "-m", "torch.distributed.run", "--standalone",
               f"--nproc_per_node={GPUS}", str(Path(__file__).resolve()),
               str(execution_path), "--worker"]
    run_logged(command, run_dir / "train.log", env)
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, ValueError, RuntimeError, ImportError, subprocess.CalledProcessError) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        raise SystemExit(1)