#!/usr/bin/env python3
"""LoRA/QLoRA SFT runner for the Qwen CyberGym project.

This script is intentionally framework-light: it uses Transformers Trainer plus
PEFT, and consumes the YAML contracts in training/configs.
"""

from __future__ import annotations

import argparse
import inspect
import json
import os
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
import yaml


ASSISTANT_START = "<|im_start|>assistant\n"
IM_END = "<|im_end|>"


@dataclass
class TokenizedExample:
    input_ids: list[int]
    attention_mask: list[int]
    labels: list[int]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, help="YAML config under training/configs.")
    parser.add_argument(
        "--allow-missing-cybergym-baseline",
        action="store_true",
        help="Dry-run escape hatch. Real training should not use this.",
    )
    parser.add_argument("--dry-run", action="store_true", help="Validate config/data/model class imports without training.")
    return parser.parse_args()


def read_yaml(path: str | Path) -> dict[str, Any]:
    with Path(path).open("r", encoding="utf-8") as fh:
        payload = yaml.safe_load(fh) or {}
    if not isinstance(payload, dict):
        raise TypeError(f"Expected a YAML mapping in {path}")
    return payload


def run_gate_check(config_path: str, allow_missing: bool) -> None:
    cmd = [sys.executable, "training/scripts/check_training_gates.py", "--config", config_path]
    if allow_missing:
        cmd.append("--allow-missing-cybergym-baseline")
    subprocess.run(cmd, check=True)


def import_training_deps():
    try:
        import transformers
        from datasets import Dataset
        from peft import LoraConfig, TaskType, get_peft_model
        from transformers import (
            AutoTokenizer,
            BitsAndBytesConfig,
            Trainer,
            TrainingArguments,
        )
    except Exception as exc:  # pragma: no cover - exercised on remote host
        raise RuntimeError(f"Missing training dependency: {exc!r}") from exc

    return {
        "transformers": transformers,
        "Dataset": Dataset,
        "AutoTokenizer": AutoTokenizer,
        "BitsAndBytesConfig": BitsAndBytesConfig,
        "Trainer": Trainer,
        "TrainingArguments": TrainingArguments,
        "LoraConfig": LoraConfig,
        "TaskType": TaskType,
        "get_peft_model": get_peft_model,
    }


def torch_dtype(name: str):
    if name == "auto":
        return "auto"
    return {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }[name]


def load_model(transformers_module, model_cfg: dict[str, Any], quantization_cfg: dict[str, Any] | None):
    model_name = model_cfg["name_or_path"]
    dtype = torch_dtype(str(model_cfg.get("dtype", "bfloat16")))
    kwargs: dict[str, Any] = {
        "trust_remote_code": bool(model_cfg.get("trust_remote_code", True)),
        "torch_dtype": dtype,
    }
    if quantization_cfg:
        deps = import_training_deps()
        kwargs["quantization_config"] = deps["BitsAndBytesConfig"](**quantization_cfg)
        kwargs["device_map"] = model_cfg.get("device_map", "auto")

    attn = model_cfg.get("attn_implementation")
    if attn:
        kwargs["attn_implementation"] = attn

    candidate_class_names = [
        "AutoModelForMultimodalLM",
        "AutoModelForImageTextToText",
        "AutoModelForVision2Seq",
        "AutoModelForCausalLM",
    ]
    errors: list[str] = []
    for class_name in candidate_class_names:
        model_cls = getattr(transformers_module, class_name, None)
        if model_cls is None:
            errors.append(f"{class_name}: not available")
            continue
        try:
            return model_cls.from_pretrained(model_name, **kwargs)
        except Exception as exc:
            errors.append(f"{class_name}: {exc!r}")

    fallback_attn = model_cfg.get("fallback_attn_implementation")
    if fallback_attn and attn and fallback_attn != attn:
        kwargs["attn_implementation"] = fallback_attn
        for class_name in candidate_class_names:
            model_cls = getattr(transformers_module, class_name, None)
            if model_cls is None:
                continue
            try:
                return model_cls.from_pretrained(model_name, **kwargs)
            except Exception as exc:
                errors.append(f"{class_name} with fallback attn: {exc!r}")

    raise RuntimeError("Could not load model:\n" + "\n".join(errors))


def freeze_by_name(model, patterns: list[str]) -> int:
    frozen = 0
    lowered = [pattern.lower() for pattern in patterns]
    for name, param in model.named_parameters():
        if any(pattern in name.lower() for pattern in lowered):
            param.requires_grad = False
            frozen += param.numel()
    return frozen


def build_lora_config(lora_cls, task_type, lora_cfg: dict[str, Any]):
    payload: dict[str, Any] = {
        "task_type": task_type.CAUSAL_LM,
        "r": int(lora_cfg["r"]),
        "lora_alpha": int(lora_cfg["alpha"]),
        "lora_dropout": float(lora_cfg.get("dropout", 0.0)),
        "bias": "none",
    }
    target_modules = lora_cfg.get("target_modules", "all-linear")
    payload["target_modules"] = target_modules

    if lora_cfg.get("use_rslora") is not None:
        payload["use_rslora"] = bool(lora_cfg["use_rslora"])

    signature = inspect.signature(lora_cls)
    if "exclude_modules" in signature.parameters and lora_cfg.get("exclude_modules"):
        payload["exclude_modules"] = lora_cfg["exclude_modules"]

    accepted = {key: value for key, value in payload.items() if key in signature.parameters}
    return lora_cls(**accepted)


def read_jsonl(path: str | Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with Path(path).open("r", encoding="utf-8") as fh:
        for line_no, line in enumerate(fh, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                payload = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSON in {path}:{line_no}: {exc}") from exc
            rows.append(payload)
    return rows


def validate_think_blocks(row: dict[str, Any], require: bool) -> None:
    if not require or "messages" not in row:
        return
    for message in row["messages"]:
        if message.get("role") == "assistant":
            content = str(message.get("content", ""))
            if "<think>" not in content or "</think>" not in content:
                row_id = row.get("id", "<unknown>")
                raise ValueError(f"Assistant message missing <think> block in row {row_id}")


def assistant_char_mask(rendered: str) -> list[bool]:
    mask = [False] * len(rendered)
    cursor = 0
    while True:
        start = rendered.find(ASSISTANT_START, cursor)
        if start == -1:
            break
        content_start = start + len(ASSISTANT_START)
        end = rendered.find(IM_END, content_start)
        if end == -1:
            end = len(rendered)
        for idx in range(content_start, end):
            mask[idx] = True
        cursor = end + len(IM_END)
    return mask


def tokenize_row(tokenizer, row: dict[str, Any], max_seq_length: int, require_think: bool) -> TokenizedExample:
    validate_think_blocks(row, require_think)

    if "messages" in row:
        try:
            rendered = tokenizer.apply_chat_template(
                row["messages"],
                tokenize=False,
                add_generation_prompt=False,
                preserve_thinking=True,
            )
        except TypeError:
            rendered = tokenizer.apply_chat_template(
                row["messages"],
                tokenize=False,
                add_generation_prompt=False,
            )
        mask = assistant_char_mask(rendered)
    elif "text" in row:
        rendered = str(row["text"])
        mask = [True] * len(rendered)
    else:
        raise ValueError("Each row needs either messages or text")

    encoded = tokenizer(
        rendered,
        add_special_tokens=False,
        truncation=True,
        max_length=max_seq_length,
        return_offsets_mapping=True,
    )

    labels: list[int] = []
    for token_id, (start, end) in zip(encoded["input_ids"], encoded["offset_mapping"], strict=True):
        if end <= start:
            labels.append(-100)
            continue
        supervised = any(mask[idx] for idx in range(start, min(end, len(mask))))
        labels.append(token_id if supervised else -100)

    return TokenizedExample(
        input_ids=list(encoded["input_ids"]),
        attention_mask=[1] * len(encoded["input_ids"]),
        labels=labels,
    )


def pack_examples(examples: list[TokenizedExample], max_seq_length: int) -> list[TokenizedExample]:
    packed: list[TokenizedExample] = []
    cur_ids: list[int] = []
    cur_labels: list[int] = []

    def flush() -> None:
        nonlocal cur_ids, cur_labels
        if cur_ids:
            packed.append(TokenizedExample(cur_ids, [1] * len(cur_ids), cur_labels))
            cur_ids = []
            cur_labels = []

    for example in examples:
        ids = example.input_ids
        labels = example.labels
        if len(ids) > max_seq_length:
            ids = ids[:max_seq_length]
            labels = labels[:max_seq_length]
        if cur_ids and len(cur_ids) + len(ids) > max_seq_length:
            flush()
        if len(ids) == max_seq_length:
            packed.append(TokenizedExample(ids, [1] * len(ids), labels))
        else:
            cur_ids.extend(ids)
            cur_labels.extend(labels)
    flush()
    return packed


class CausalCollator:
    def __init__(self, pad_token_id: int, label_pad_token_id: int = -100):
        self.pad_token_id = pad_token_id
        self.label_pad_token_id = label_pad_token_id

    def __call__(self, features: list[dict[str, list[int]]]) -> dict[str, torch.Tensor]:
        max_len = max(len(feature["input_ids"]) for feature in features)
        input_ids = []
        attention_mask = []
        labels = []
        for feature in features:
            pad = max_len - len(feature["input_ids"])
            input_ids.append(feature["input_ids"] + [self.pad_token_id] * pad)
            attention_mask.append(feature["attention_mask"] + [0] * pad)
            labels.append(feature["labels"] + [self.label_pad_token_id] * pad)
        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }


def make_dataset(dataset_cls, rows: list[dict[str, Any]], tokenizer, data_cfg: dict[str, Any]):
    max_seq_length = int(data_cfg["max_seq_length"])
    require_think = bool(data_cfg.get("require_think_blocks", False))
    tokenized = [tokenize_row(tokenizer, row, max_seq_length, require_think) for row in rows]
    if data_cfg.get("packing", False):
        tokenized = pack_examples(tokenized, max_seq_length)
    payload = [
        {"input_ids": item.input_ids, "attention_mask": item.attention_mask, "labels": item.labels}
        for item in tokenized
    ]
    return dataset_cls.from_list(payload)


def training_args_kwargs(training_arguments_cls, run_cfg: dict[str, Any], training_cfg: dict[str, Any]) -> dict[str, Any]:
    output_dir = run_cfg["output_dir"]
    payload: dict[str, Any] = {
        "output_dir": output_dir,
        "overwrite_output_dir": False,
        "learning_rate": float(training_cfg["learning_rate"]),
        "lr_scheduler_type": training_cfg.get("lr_scheduler_type", "cosine"),
        "warmup_ratio": float(training_cfg.get("warmup_ratio", 0.03)),
        "num_train_epochs": float(training_cfg["num_train_epochs"]),
        "per_device_train_batch_size": int(training_cfg["per_device_train_batch_size"]),
        "per_device_eval_batch_size": int(training_cfg.get("per_device_eval_batch_size", 1)),
        "gradient_accumulation_steps": int(training_cfg.get("gradient_accumulation_steps", 1)),
        "gradient_checkpointing": bool(training_cfg.get("gradient_checkpointing", True)),
        "max_grad_norm": float(training_cfg.get("max_grad_norm", 1.0)),
        "logging_steps": int(training_cfg.get("logging_steps", 10)),
        "save_strategy": training_cfg.get("save_strategy", "steps"),
        "bf16": bool(training_cfg.get("bf16", True)),
        "tf32": bool(training_cfg.get("tf32", True)),
        "report_to": ["wandb"] if os.getenv("WANDB_API_KEY") else [],
        "remove_unused_columns": False,
        "seed": int(run_cfg.get("seed", 1337)),
    }

    if "max_steps" in training_cfg:
        payload["max_steps"] = int(training_cfg["max_steps"])
    if "save_steps" in training_cfg:
        payload["save_steps"] = int(training_cfg["save_steps"])
    if "eval_steps" in training_cfg:
        payload["eval_steps"] = int(training_cfg["eval_steps"])
    if "eval_strategy" in training_cfg:
        payload["eval_strategy"] = training_cfg["eval_strategy"]
    elif "evaluation_strategy" in training_cfg:
        payload["evaluation_strategy"] = training_cfg["evaluation_strategy"]
    elif "eval_steps" in training_cfg:
        payload["eval_strategy"] = "steps"
    else:
        payload["eval_strategy"] = "epoch"

    signature = inspect.signature(training_arguments_cls)
    return {key: value for key, value in payload.items() if key in signature.parameters}


def build_callbacks(run_cfg: dict[str, Any], config: dict[str, Any]):
    """Metrics logger (always) + optional in-training held-out benchmark."""
    from transformers import TrainerCallback

    output_dir = Path(run_cfg["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    metrics_path = output_dir / "metrics.jsonl"
    progress_path = output_dir / "eval_progress.jsonl"

    class JsonlMetricsCallback(TrainerCallback):
        """Append every Trainer log line to metrics.jsonl for the watcher."""

        def on_log(self, args, state, control, logs=None, **kwargs):
            if not logs or not state.is_world_process_zero:
                return
            row = {k: v for k, v in logs.items() if isinstance(v, (int, float))}
            row.update({"step": state.global_step, "epoch": state.epoch, "ts": time.time()})
            with metrics_path.open("a", encoding="utf-8") as fh:
                fh.write(json.dumps(row) + "\n")

    callbacks = [JsonlMetricsCallback()]

    eval_cfg = config.get("in_training_eval") or {}
    if not eval_cfg.get("enabled"):
        return callbacks

    eval_files = eval_cfg.get("eval_files", [])
    sample = int(eval_cfg.get("sample_per_set", 60))
    max_new = int(eval_cfg.get("max_new_tokens", 256))
    enable_thinking = bool(eval_cfg.get("enable_thinking", False))
    eval_every_steps = int(eval_cfg.get("eval_every_steps", 0))
    base_acc: dict[str, float] = {}
    base_json = eval_cfg.get("base_eval_json")
    if base_json and Path(base_json).is_file():
        try:
            payload = json.loads(Path(base_json).read_text(encoding="utf-8"))
            base_acc = {r["kind"]: r["accuracy"] for r in payload.get("results", []) if "kind" in r}
        except Exception:
            base_acc = {}

    class PeriodicEvalCallback(TrainerCallback):
        """Run the held-out benchmark on the training model at a step interval + each epoch."""

        def _run(self, state, kwargs):
            if not state.is_world_process_zero:
                return
            model = kwargs.get("model")
            tokenizer = kwargs.get("processing_class") or kwargs.get("tokenizer")
            if model is None or tokenizer is None:
                return
            sys.path.insert(0, str(Path(__file__).resolve().parent))
            try:
                from intraining_eval import run_eval_sets

                sets = run_eval_sets(model, tokenizer, eval_files, sample, max_new, enable_thinking)
                deltas = {}
                for metrics in sets.values():
                    kind = metrics.get("kind")
                    if kind in base_acc and "accuracy" in metrics:
                        deltas[kind] = round(metrics["accuracy"] - base_acc[kind], 4)
                row = {
                    "epoch": state.epoch, "step": state.global_step, "ts": time.time(),
                    "sets": sets, "base": base_acc, "deltas_vs_base": deltas,
                }
            except Exception as exc:  # never let a benchmark kill training
                row = {"epoch": state.epoch, "step": state.global_step, "ts": time.time(),
                       "error": repr(exc)}
            with progress_path.open("a", encoding="utf-8") as fh:
                fh.write(json.dumps(row) + "\n")
            print(f"[in-training-eval] step={state.global_step} epoch={state.epoch}: "
                  f"{row.get('deltas_vs_base', row.get('error'))}")

        def on_step_end(self, args, state, control, **kwargs):
            if eval_every_steps and state.global_step > 0 and state.global_step % eval_every_steps == 0:
                self._run(state, kwargs)

        def on_epoch_end(self, args, state, control, **kwargs):
            self._run(state, kwargs)

    callbacks.append(PeriodicEvalCallback())
    return callbacks


def main() -> int:
    args = parse_args()
    run_gate_check(args.config, args.allow_missing_cybergym_baseline)

    config = read_yaml(args.config)
    deps = import_training_deps()

    run_cfg = config["run"]
    model_cfg = config["model"]
    data_cfg = config["data"]
    training_cfg = config["training"]

    tokenizer = deps["AutoTokenizer"].from_pretrained(
        model_cfg["name_or_path"],
        trust_remote_code=bool(model_cfg.get("trust_remote_code", True)),
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    train_rows = read_jsonl(data_cfg["train_jsonl"])
    val_rows = read_jsonl(data_cfg["validation_jsonl"])
    train_ds = make_dataset(deps["Dataset"], train_rows, tokenizer, data_cfg)
    eval_ds = make_dataset(deps["Dataset"], val_rows, tokenizer, data_cfg)

    if args.dry_run:
        print(f"Dry run ok: train={len(train_ds)} eval={len(eval_ds)}")
        return 0

    quantization_cfg = training_cfg.get("quantization") if training_cfg.get("method") == "qlora" else None
    model = load_model(deps["transformers"], model_cfg, quantization_cfg)

    freeze_patterns: list[str] = []
    if model_cfg.get("freeze_vision_tower", True):
        freeze_patterns.extend(["visual", "vision_tower", "multi_modal_projector"])
    if model_cfg.get("freeze_mtp_head", True):
        freeze_patterns.extend(["mtp"])
    frozen_params = freeze_by_name(model, freeze_patterns)
    print(f"Frozen parameter elements by name pattern: {frozen_params}")

    lora_config = build_lora_config(deps["LoraConfig"], deps["TaskType"], training_cfg["lora"])
    model = deps["get_peft_model"](model, lora_config)
    model.print_trainable_parameters()

    if training_cfg.get("gradient_checkpointing", True):
        model.config.use_cache = False

    training_args = deps["TrainingArguments"](
        **training_args_kwargs(deps["TrainingArguments"], run_cfg, training_cfg)
    )
    trainer_kwargs = {
        "model": model,
        "args": training_args,
        "train_dataset": train_ds,
        "eval_dataset": eval_ds,
        "data_collator": CausalCollator(tokenizer.pad_token_id),
    }
    trainer_signature = inspect.signature(deps["Trainer"])
    if "processing_class" in trainer_signature.parameters:
        trainer_kwargs["processing_class"] = tokenizer
    elif "tokenizer" in trainer_signature.parameters:
        trainer_kwargs["tokenizer"] = tokenizer
    trainer_kwargs["callbacks"] = build_callbacks(run_cfg, config)
    trainer = deps["Trainer"](**trainer_kwargs)
    trainer.train()
    trainer.save_model(run_cfg["output_dir"])
    tokenizer.save_pretrained(run_cfg["output_dir"])
    return 0


if __name__ == "__main__":
    raise SystemExit(main())