#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Unified training script: SFT (QLoRA + masked loss) and DPO (TRL) in ONE file.

UX requirement (as you asked):
- You run:   python3 train_unified.py
- It shows a simple menu 1/2/3
- You press 1, 2, or 3
- Script relaunches itself with accelerate (2 GPUs) and runs the chosen stage
  using DEFAULT paths/hparams already embedded.

Notes:
- Menu is shown ONLY once (before accelerate spawns processes).
- After relaunch, each process runs normally with LOCAL_RANK set by accelerate.
"""

import os
import sys
import json
import math
import argparse
import subprocess
import inspect
import traceback
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

import torch

from datasets import load_dataset, load_from_disk
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    Trainer as HFTrainer,
    TrainingArguments,
)
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    PeftModel,
)
from trl import DPOTrainer, DPOConfig


# --------------------------
# Defaults (same spirit as your original scripts)
# --------------------------
DEFAULT_MODEL_DIR = "/home/hyrenko/Diploma/models/llama3.1-8b"

# SFT: source -> jsonl cache
SFT_SRC_DATASET_ID = "PKU-Alignment/PKU-SafeRLHF-30K"
SFT_SRC_SPLIT = "train"
SFT_JSONL_PATH = "./data/pku_sft.jsonl"

# Outputs
DEFAULT_SFT_OUT = "./out/llama3_1_8b_sft_safety_lora_masked"
DEFAULT_DPO_OUT = "./out/llama3_1_8b_dpo_safety_lora"
DEFAULT_DPO_DATA = "./data/pku_dpo"

# Prompt wrappers (keep consistent)
INSTR_PREFIX = "### Instruction:\n"
RESP_PREFIX = "\n\n### Response:\n"


# --------------------------
# Simple interactive menu (shown only before accelerate)
# --------------------------
def choose_stage_simple() -> str:
    """
    Returns: 'sft' | 'dpo' | 'both'
    """
    print("\n=== TRAIN MENU ===")
    print("1) SFT (QLoRA masked-loss)")
    print("2) DPO (over SFT adapter)")
    print("3) BOTH (SFT -> DPO)")
    while True:
        x = input("Choose 1/2/3: ").strip()
        if x == "1":
            return "sft"
        if x == "2":
            return "dpo"
        if x == "3":
            return "both"
        print("Only 1, 2, or 3.")


def relaunch_with_accelerate(num_processes: int, extra_argv: List[str]):
    """
    Relaunches the script via accelerate (2 processes by default).
    Only runs if NOT already inside accelerate (LOCAL_RANK not set).
    """
    if os.environ.get("LOCAL_RANK") is not None:
        return

    cmd = [
    sys.executable,
    "-m",
    "accelerate.commands.launch",
    "--num_processes",
    str(num_processes),
    "--mixed_precision",
    "fp16",
    "--main_process_port",
    "0",
    sys.argv[0],
    *extra_argv,
    ]
    print("\n[AUTO] Relaunching with accelerate:\n" + " ".join(cmd) + "\n")
    raise SystemExit(subprocess.call(cmd))


# --------------------------
# DDP helpers
# --------------------------
def ddp_info():
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    rank = int(os.environ.get("RANK", "0"))
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    return local_rank, rank, world_size


def ddp_barrier():
    if torch.distributed.is_available() and torch.distributed.is_initialized():
        torch.distributed.barrier()


# --------------------------
# TRL <-> Transformers patches (to avoid signature mismatches)
# --------------------------
def patch_trl_dpo_trainer_if_needed():
    # Patch get_batch_samples
    try:
        sig = list(inspect.signature(DPOTrainer.get_batch_samples).parameters.keys())
        if len(sig) >= 3 and sig[1] == "model":
            if not hasattr(DPOTrainer, "get_batch_samples_trl"):
                DPOTrainer.get_batch_samples_trl = DPOTrainer.get_batch_samples
            DPOTrainer.get_batch_samples = HFTrainer.get_batch_samples
            print("[PATCH] Patched DPOTrainer.get_batch_samples -> HF Trainer.get_batch_samples")
    except Exception:
        print("[PATCH] Could not inspect/patch DPOTrainer.get_batch_samples")

    # Patch compute_loss
    try:
        sig = inspect.signature(DPOTrainer.compute_loss)
        if "num_items_in_batch" not in sig.parameters:
            if not hasattr(DPOTrainer, "compute_loss_trl"):
                DPOTrainer.compute_loss_trl = DPOTrainer.compute_loss

            def compute_loss_patched(self, model, inputs, return_outputs=False, num_items_in_batch=None, **kwargs):
                orig = self.compute_loss_trl
                osig = inspect.signature(orig)
                if "return_outputs" in osig.parameters:
                    return orig(model, inputs, return_outputs=return_outputs)
                out = orig(model, inputs)
                if return_outputs:
                    return out, None
                return out

            DPOTrainer.compute_loss = compute_loss_patched
            print("[PATCH] Patched DPOTrainer.compute_loss to accept num_items_in_batch")
    except Exception:
        print("[PATCH] Could not inspect/patch DPOTrainer.compute_loss")


def assert_is_model(obj, name="model"):
    if not hasattr(obj, "generate"):
        raise TypeError(
            "[FATAL] {} is not a HF CausalLM model. type={}; missing .generate().".format(name, type(obj))
        )


# --------------------------
# SFT helpers
# --------------------------
def wrap_prompt(p: str) -> str:
    p = (p or "").strip()
    return "{}{}{}".format(INSTR_PREFIX, p, RESP_PREFIX)


def ensure_sft_jsonl_exists(rank: int):
    """
    Creates ./data/pku_sft.jsonl from PKU-SafeRLHF-30K if missing.
    Only rank0 writes; others wait.
    """
    os.makedirs("./data", exist_ok=True)

    if os.path.exists(SFT_JSONL_PATH) and os.path.getsize(SFT_JSONL_PATH) > 0:
        if rank == 0:
            print("[INFO] Found existing SFT jsonl:", SFT_JSONL_PATH)
        ddp_barrier()
        return

    if rank == 0:
        print("[INFO] SFT jsonl not found. Building from {} ...".format(SFT_SRC_DATASET_ID))
        ds = load_dataset(SFT_SRC_DATASET_ID, split=SFT_SRC_SPLIT)

        rows = 0
        dropped = 0

        with open(SFT_JSONL_PATH, "w", encoding="utf-8") as f:
            for ex in ds:
                prompt = (ex.get("prompt") or "").strip()
                r0 = (ex.get("response_0") or "").strip()
                r1 = (ex.get("response_1") or "").strip()
                sid = ex.get("safer_response_id", None)

                if not prompt or not r0 or not r1 or sid not in (0, 1):
                    dropped += 1
                    continue

                safe = r0 if sid == 0 else r1
                obj = {"prompt": wrap_prompt(prompt), "response": safe}
                f.write(json.dumps(obj, ensure_ascii=False) + "\n")
                rows += 1

        print("[OK] Built SFT jsonl:", SFT_JSONL_PATH)
        print("[OK] rows={}, dropped={}".format(rows, dropped))

    ddp_barrier()


@dataclass
class DataCollatorForCausalLMWithLabels:
    tokenizer: Any

    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
        pad_batch = self.tokenizer.pad(
            {
                "input_ids": [f["input_ids"] for f in features],
                "attention_mask": [f["attention_mask"] for f in features],
            },
            padding=True,
            return_tensors="pt",
        )

        max_len = pad_batch["input_ids"].shape[1]
        labels = []
        for f in features:
            lab = f["labels"]
            if len(lab) < max_len:
                lab = lab + [-100] * (max_len - len(lab))
            else:
                lab = lab[:max_len]
            labels.append(lab)

        pad_batch["labels"] = torch.tensor(labels, dtype=torch.long)
        return pad_batch


def build_quant_base(model_dir: str, device_map: Dict[str, int]):
    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
    )
    base = AutoModelForCausalLM.from_pretrained(
        model_dir,
        quantization_config=bnb,
        device_map=device_map,
        torch_dtype=torch.float16,
    )
    base.config.use_cache = False
    return base


def run_sft(
    model_dir: str,
    out_dir: str,
    epochs: float,
    lr: float,
    bs: int,
    gas: int,
    seq: int,
    lora_r: int,
    device_map: Dict[str, int],
    rank: int,
    world_size: int,
) -> str:
    """
    Runs SFT and returns the path to the saved LoRA adapter (out_dir).
    """
    os.makedirs(out_dir, exist_ok=True)

    ensure_sft_jsonl_exists(rank)

    ds = load_dataset("json", data_files=SFT_JSONL_PATH, split="train")

    tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    base = build_quant_base(model_dir, device_map)
    base = prepare_model_for_kbit_training(base)

    lora = LoraConfig(
        r=lora_r,
        lora_alpha=32,
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "v_proj"],
    )
    model = get_peft_model(base, lora)

    def tokenize_and_mask(example):
        prompt = example["prompt"]
        response = example["response"]

        p = tokenizer(prompt, add_special_tokens=False)
        r = tokenizer(response + tokenizer.eos_token, add_special_tokens=False)

        input_ids = p["input_ids"] + r["input_ids"]
        attention_mask = [1] * len(input_ids)
        labels = [-100] * len(p["input_ids"]) + r["input_ids"]

        if len(input_ids) > seq:
            input_ids = input_ids[:seq]
            attention_mask = attention_mask[:seq]
            labels = labels[:seq]

        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

    ds_tok = ds.map(tokenize_and_mask, remove_columns=ds.column_names)

    if rank == 0:
        total_examples = len(ds_tok)
        effective_bs = bs * world_size * gas
        steps_per_epoch = int(math.ceil(float(total_examples) / float(effective_bs)))
        print("[SFT] base={}".format(model_dir))
        print("[SFT] out={}".format(out_dir))
        print("[SFT] examples={}, world_size={}".format(total_examples, world_size))
        print("[SFT] per_device_bs={}, gas={}, effective_batch={}".format(bs, gas, effective_bs))
        print("[SFT] approx steps/epoch ≈ {}".format(steps_per_epoch))

    collator = DataCollatorForCausalLMWithLabels(tokenizer)

    train_args = TrainingArguments(
        output_dir=out_dir,
        num_train_epochs=epochs,
        learning_rate=lr,
        per_device_train_batch_size=bs,
        gradient_accumulation_steps=gas,
        logging_steps=10,
        save_steps=200,
        save_total_limit=2,
        fp16=True,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        ddp_find_unused_parameters=False,
        report_to="none",
    )

    trainer = HFTrainer(
        model=model,
        args=train_args,
        train_dataset=ds_tok,
        data_collator=collator,
        tokenizer=tokenizer,
    )

    trainer.train()

    # Save only on rank0, then barrier so others wait
    if rank == 0:
        model.save_pretrained(out_dir)
        tokenizer.save_pretrained(out_dir)
        print("[OK] SFT LoRA adapter saved to:", out_dir)
    ddp_barrier()

    return out_dir


# --------------------------
# DPO stage
# --------------------------
def run_dpo(
    model_dir: str,
    sft_adapter: str,
    dpo_data_dir: str,
    out_dir: str,
    epochs: float,
    lr: float,
    beta: float,
    bs: int,
    gas: int,
    max_prompt: int,
    max_len: int,
    device_map: Dict[str, int],
    rank: int,
):
    patch_trl_dpo_trainer_if_needed()
    os.makedirs(out_dir, exist_ok=True)

    ds = load_from_disk(dpo_data_dir)

    tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    base_model = build_quant_base(model_dir, device_map)
    assert_is_model(base_model, "base_model")

    # Load SFT adapter and keep training it with DPO
    model = PeftModel.from_pretrained(base_model, sft_adapter, is_trainable=True)
    assert_is_model(model, "model(peft)")

    if rank == 0:
        print("[DPO] base:", model_dir)
        print("[DPO] sft_adapter:", sft_adapter)
        print("[DPO] dpo_data:", dpo_data_dir)
        print("[DPO] out:", out_dir)

    dpo_args = DPOConfig(
        output_dir=out_dir,
        num_train_epochs=epochs,
        learning_rate=lr,
        per_device_train_batch_size=bs,
        gradient_accumulation_steps=gas,
        logging_steps=10,
        save_steps=200,
        save_total_limit=2,
        fp16=True,
        report_to="none",
        beta=beta,
        max_prompt_length=max_prompt,
        max_length=max_len,
        gradient_checkpointing=True,
        ddp_find_unused_parameters=False,
    )

    trainer = DPOTrainer(
        model=model,
        args=dpo_args,
        train_dataset=ds,
        tokenizer=tokenizer,
    )

    trainer.train()
    trainer.save_model(out_dir)

    if rank == 0:
        print("[OK] DPO LoRA adapter saved to:", out_dir)


# --------------------------
# Main
# --------------------------
def main():
    ap = argparse.ArgumentParser(description="Unified SFT(QLoRA masked) + DPO training script")

    # Stage (menu will override this in the initial run)
    ap.add_argument("--stage", choices=["sft", "dpo", "both"], default="both")

    ap.add_argument("--model", default=DEFAULT_MODEL_DIR)

    # SFT args
    ap.add_argument("--sft_out", default=DEFAULT_SFT_OUT)
    ap.add_argument("--sft_epochs", type=float, default=1.0)
    ap.add_argument("--sft_lr", type=float, default=2e-4)
    ap.add_argument("--sft_bs", type=int, default=1)
    ap.add_argument("--sft_gas", type=int, default=16)
    ap.add_argument("--sft_seq", type=int, default=1024)
    ap.add_argument("--lora_r", type=int, default=16)

    # DPO args
    ap.add_argument("--sft_adapter", default=None, help="If omitted, uses --sft_out")
    ap.add_argument("--dpo_data", default=DEFAULT_DPO_DATA)
    ap.add_argument("--dpo_out", default=DEFAULT_DPO_OUT)
    ap.add_argument("--dpo_epochs", type=float, default=1.0)
    ap.add_argument("--dpo_lr", type=float, default=5e-6)
    ap.add_argument("--beta", type=float, default=0.1)
    ap.add_argument("--dpo_bs", type=int, default=1)
    ap.add_argument("--dpo_gas", type=int, default=16)
    ap.add_argument("--max_prompt", type=int, default=512)
    ap.add_argument("--max_len", type=int, default=1024)

    args = ap.parse_args()

    # --------------------------
    # AUTO MENU (only once, BEFORE accelerate spawns processes)
    # --------------------------
    if os.environ.get("LOCAL_RANK") is None and os.environ.get("AUTO_ACCELERATE_RELAUNCH") != "1":
        os.environ["AUTO_ACCELERATE_RELAUNCH"] = "1"  # prevent menu duplication

        stage = choose_stage_simple()

        # Build argv using defaults from args (you only press 1/2/3)
        argv = [
            "--stage",
            stage,
            "--model",
            args.model,

            "--sft_out",
            args.sft_out,
            "--sft_epochs",
            str(args.sft_epochs),
            "--sft_lr",
            str(args.sft_lr),
            "--sft_bs",
            str(args.sft_bs),
            "--sft_gas",
            str(args.sft_gas),
            "--sft_seq",
            str(args.sft_seq),
            "--lora_r",
            str(args.lora_r),

            "--dpo_data",
            args.dpo_data,
            "--dpo_out",
            args.dpo_out,
            "--dpo_epochs",
            str(args.dpo_epochs),
            "--dpo_lr",
            str(args.dpo_lr),
            "--beta",
            str(args.beta),
            "--dpo_bs",
            str(args.dpo_bs),
            "--dpo_gas",
            str(args.dpo_gas),
            "--max_prompt",
            str(args.max_prompt),
            "--max_len",
            str(args.max_len),
        ]

        if stage in ("dpo", "both"):
            sft_adapter = args.sft_adapter or args.sft_out
            argv += ["--sft_adapter", sft_adapter]

        # Relaunch under accelerate (2 GPUs)
        relaunch_with_accelerate(num_processes=2, extra_argv=argv)

    # --------------------------
    # We are inside accelerate now
    # --------------------------
    local_rank, rank, world_size = ddp_info()
    torch.cuda.set_device(local_rank)
    device_map = {"": torch.cuda.current_device()}

    # Deterministic seed (as in your DPO script style)
    torch.manual_seed(42)

    sft_adapter_path = args.sft_adapter or args.sft_out

    if args.stage in ("sft", "both"):
        sft_adapter_path = run_sft(
            model_dir=args.model,
            out_dir=args.sft_out,
            epochs=args.sft_epochs,
            lr=args.sft_lr,
            bs=args.sft_bs,
            gas=args.sft_gas,
            seq=args.sft_seq,
            lora_r=args.lora_r,
            device_map=device_map,
            rank=rank,
            world_size=world_size,
        )

    if args.stage in ("dpo", "both"):
        # If "both", barrier already happened after saving SFT adapter
        run_dpo(
            model_dir=args.model,
            sft_adapter=sft_adapter_path,
            dpo_data_dir=args.dpo_data,
            out_dir=args.dpo_out,
            epochs=args.dpo_epochs,
            lr=args.dpo_lr,
            beta=args.beta,
            bs=args.dpo_bs,
            gas=args.dpo_gas,
            max_prompt=args.max_prompt,
            max_len=args.max_len,
            device_map=device_map,
            rank=rank,
        )


if __name__ == "__main__":
    try:
        main()
    except Exception:
        r = os.environ.get("RANK", "?")
        lr = os.environ.get("LOCAL_RANK", "?")
        print("\n[FATAL] Exception on RANK={} LOCAL_RANK={}\n".format(r, lr))
        print(traceback.format_exc())
        raise