"""RLHF (PPO) trainer — TRL PPOTrainer wrapper.

Online preference optimization with a learned reward model. Less efficient
than DPO for offline pairs; kept here as the canonical surface for online
RLHF runs.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

from mindxtrain.config.schema import XTrainConfig


def run_rlhf(cfg: XTrainConfig, out_dir: Path) -> Path:
    """Run a PPO/RLHF fine-tune; return the checkpoint directory."""
    try:
        from datasets import load_dataset
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from trl import PPOConfig, PPOTrainer
    except ImportError as exc:
        msg = "TRL stack not installed; run `uv sync --extra ml`."
        raise RuntimeError(msg) from exc

    method: Any = cfg.train.method
    if getattr(method, "kind", "") != "rlhf":
        msg = f"run_rlhf expects train.method.kind == 'rlhf'; got {getattr(method, 'kind', None)!r}"
        raise ValueError(msg)

    reward_model_path = getattr(method, "reward_model_path", "")
    if not reward_model_path:
        msg = "rlhf requires train.method.reward_model_path"
        raise ValueError(msg)

    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(cfg.model.name)
    policy = AutoModelForCausalLM.from_pretrained(cfg.model.name)
    ref_policy = AutoModelForCausalLM.from_pretrained(cfg.model.name)
    reward_model = AutoModelForCausalLM.from_pretrained(reward_model_path)

    train_ds = load_dataset(cfg.data.hf_id, split=getattr(cfg.data, "split", "train"))

    ppo_cfg = PPOConfig(
        output_dir=str(out_dir),
        learning_rate=cfg.train.optim.learning_rate,
        per_device_train_batch_size=cfg.train.micro_batch_size,
        gradient_accumulation_steps=cfg.train.gradient_accumulation_steps,
        num_train_epochs=cfg.train.num_epochs,
    )
    trainer = PPOTrainer(
        args=ppo_cfg,
        model=policy,
        ref_model=ref_policy,
        reward_model=reward_model,
        train_dataset=train_ds,
        processing_class=tokenizer,
    )
    trainer.train()
    trainer.save_model(str(out_dir))
    return out_dir


__all__ = ["run_rlhf"]