"""Fine-tune a causal language model to produce a LoRA or aLoRA adapter.

Loads a JSONL dataset of ``item``/``label`` pairs, applies an 80/20 train/validation
split, and trains using HuggingFace PEFT and TRL's ``SFTTrainer`` — saving the
checkpoint with the lowest validation loss. Supports CUDA, MPS (macOS,
PyTorch ≥ 2.8), and CPU device selection, and handles the
``alora_invocation_tokens`` configuration required for aLoRA training.
"""

import json
import os
import sys
import warnings

import torch
import typer
from datasets import Dataset
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback
from trl import DataCollatorForCompletionOnlyLM, SFTConfig, SFTTrainer

# Handle MPS with old PyTorch versions on macOS only
# Accelerate's GradScaler requires PyTorch >= 2.8.0 for MPS
if sys.platform == "darwin" and hasattr(torch.backends, "mps"):
    if torch.backends.mps.is_available():
        pytorch_version = tuple(int(x) for x in torch.__version__.split(".")[:2])
        if pytorch_version < (2, 8):
            # Disable MPS detection to force CPU usage on macOS
            # This must be done before any models or tensors are initialized
            torch.backends.mps.is_available = lambda: False  # type: ignore[assignment]
            torch.backends.mps.is_built = lambda: False  # type: ignore[assignment]
            os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "0"
            warnings.warn(
                "MPS is available but PyTorch < 2.8.0. Disabling MPS to avoid "
                "gradient scaling issues. Training will run on CPU. "
                "To use MPS, upgrade to PyTorch >= 2.8.0.",
                UserWarning,
                stacklevel=2,
            )


def load_dataset_from_json(json_path, tokenizer, invocation_prompt):
    """Load a JSONL dataset and format it for SFT training.

    Reads ``item``/``label`` pairs from a JSONL file and builds a HuggingFace
    ``Dataset`` with ``input`` and ``target`` columns. Each input is formatted as
    ``"{item}\\nRequirement: <|end_of_text|>\\n{invocation_prompt}"``.

    Args:
        json_path: Path to the JSONL file containing ``item``/``label`` pairs.
        tokenizer: HuggingFace tokenizer instance (currently unused, reserved for
            future tokenization steps).
        invocation_prompt: Invocation string appended to each input prompt.

    Returns:
        A HuggingFace ``Dataset`` with ``"input"`` and ``"target"`` string columns.
    """
    data = []
    with open(json_path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                data.append(json.loads(line))

    inputs = []
    targets = []
    for sample in data:
        item_text = sample.get("item", "")
        label_text = sample.get("label", "")
        prompt = f"{item_text}\nRequirement: <|end_of_text|>\n{invocation_prompt}"
        inputs.append(prompt)
        targets.append(label_text)
    return Dataset.from_dict({"input": inputs, "target": targets})


def formatting_prompts_func(example):
    """Concatenate input and target columns for SFT prompt formatting.

    Args:
        example: A batch dict with ``"input"`` and ``"target"`` list fields, as
            produced by HuggingFace ``Dataset.map`` in batched mode.

    Returns:
        A list of strings, each formed by concatenating the ``input`` and
        ``target`` values for a single example in the batch.
    """
    return [
        f"{example['input'][i]}{example['target'][i]}"
        for i in range(len(example["input"]))
    ]


class SaveBestModelCallback(TrainerCallback):
    """HuggingFace Trainer callback that saves the adapter at its best validation loss.

    Attributes:
        best_eval_loss (float): Lowest evaluation loss seen so far across all
            evaluation steps. Initialised to ``float("inf")``.
    """

    def __init__(self):
        self.best_eval_loss = float("inf")

    def on_evaluate(self, args, state, control, **kwargs):
        """Save the adapter weights if the current evaluation loss is a new best.

        Called automatically by the HuggingFace Trainer after each evaluation
        step. Compares the current ``eval_loss`` from ``metrics`` against
        ``best_eval_loss`` and, if lower, updates the stored best and saves the
        model to ``args.output_dir``.

        Args:
            args: ``TrainingArguments`` instance with training configuration,
                including ``output_dir``.
            state: ``TrainerState`` instance with the current training state.
            control: ``TrainerControl`` instance for controlling training flow.
            **kwargs: Additional keyword arguments provided by the Trainer,
                including ``"model"`` (the current PEFT model) and
                ``"metrics"`` (a dict containing ``"eval_loss"``).
        """
        model = kwargs["model"]
        metrics = kwargs["metrics"]
        eval_loss = metrics.get("eval_loss")
        if eval_loss is not None and eval_loss < self.best_eval_loss:
            self.best_eval_loss = eval_loss
            model.save_pretrained(args.output_dir)


class SafeSaveTrainer(SFTTrainer):
    """SFTTrainer subclass that always saves models with safe serialization enabled."""

    def save_model(self, output_dir: str | None = None, _internal_call: bool = False):
        """Save the model and tokenizer with safe serialization always enabled.

        Overrides ``SFTTrainer.save_model`` to call ``save_pretrained`` with
        ``safe_serialization=True``, ensuring weights are saved in safetensors
        format rather than the legacy pickle-based format.

        Args:
            output_dir (str | None): Directory to save the model into. If
                ``None``, the trainer's configured ``output_dir`` is used.
            _internal_call (bool): Internal flag passed through from the Trainer
                base class; not used by this override.
        """
        if self.model is not None:
            self.model.save_pretrained(output_dir, safe_serialization=True)
            if self.tokenizer is not None:
                self.tokenizer.save_pretrained(output_dir)


def train_model(
    dataset_path: str,
    base_model: str,
    output_file: str,
    prompt_file: str | None = None,
    adapter: str = "alora",
    device: str = "auto",
    run_name: str = "multiclass_run",
    epochs: int = 6,
    learning_rate: float = 6e-6,
    batch_size: int = 2,
    max_length: int = 1024,
    grad_accum: int = 4,
):
    """Fine-tune a causal language model to produce a LoRA or aLoRA adapter.

    Loads and 80/20-splits the JSONL dataset, configures PEFT with the specified
    adapter type, trains using ``SFTTrainer`` with a best-checkpoint callback, saves
    the adapter weights, and removes the PEFT-generated ``README.md`` from the output
    directory.

    Args:
        dataset_path: Path to the JSONL training dataset file.
        base_model: Hugging Face model ID or local path to the base model.
        output_file: Destination path for the trained adapter weights.
        prompt_file: Optional path to a JSON config file with an
            ``"invocation_prompt"`` key. Defaults to the aLoRA invocation token.
        adapter: Adapter type to train -- ``"alora"`` (default) or ``"lora"``.
        device: Device selection -- ``"auto"``, ``"cpu"``, ``"cuda"``, or
            ``"mps"``.
        run_name: Name of the training run (passed to ``SFTConfig``).
        epochs: Number of training epochs.
        learning_rate: Optimizer learning rate.
        batch_size: Per-device training batch size.
        max_length: Maximum token sequence length.
        grad_accum: Gradient accumulation steps.

    Raises:
        ValueError: If ``device`` is not one of ``"auto"``, ``"cpu"``,
            ``"cuda"``, or ``"mps"``.
        RuntimeError: If the GPU has insufficient VRAM to load the model
            (wraps ``NotImplementedError`` for meta tensor errors).
    """
    if prompt_file:
        # load the configurable variable invocation_prompt
        with open(prompt_file) as f:
            config = json.load(f)
        invocation_prompt = config["invocation_prompt"]
    else:
        invocation_prompt = "<|start_of_role|>check_requirement<|end_of_role|>"

    tokenizer = AutoTokenizer.from_pretrained(
        base_model, padding_side="right", trust_remote_code=True
    )
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.add_special_tokens = False

    dataset = load_dataset_from_json(dataset_path, tokenizer, invocation_prompt)
    dataset = dataset.shuffle(seed=42)
    split_idx = int(len(dataset) * 0.8)
    train_dataset = dataset.select(range(split_idx))
    val_dataset = dataset.select(range(split_idx, len(dataset)))

    if device == "auto":
        if torch.cuda.is_available():
            gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            if gpu_memory_gb < 6:
                print(
                    f"⚠️  Warning: GPU has {gpu_memory_gb:.1f}GB VRAM. "
                    "Training 3B+ models may fail. Consider using --device cpu"
                )
            device_map = "auto"
        else:
            device_map = None
    elif device == "cpu":
        device_map = None
    elif device in ["cuda", "mps"]:
        device_map = "auto"
    else:
        raise ValueError(f"Invalid device '{device}'. Use: auto, cpu, cuda, or mps")

    try:
        model_base = AutoModelForCausalLM.from_pretrained(
            base_model, device_map=device_map, use_cache=False
        )

        # `fp16=True` enables CUDA-specific mixed precision via GradScaler, which doesn't function properly on cpu or mps.
        # Check all the model's parameters to ensure it's okay to use.
        use_fp16 = all(
            param.device.type != "cpu" and param.device.type != "mps"
            for param in model_base.parameters()
        )
    except NotImplementedError as e:
        if "meta tensor" in str(e):
            raise RuntimeError(
                "Insufficient GPU memory for model. The model is too large for available VRAM. "
                "Try: (1) Use a smaller model, (2) Use a system with more GPU memory (6GB+ recommended), "
                "or (3) Train on CPU (slower but works with limited memory)."
            ) from e
        raise

    collator = DataCollatorForCompletionOnlyLM(invocation_prompt, tokenizer=tokenizer)

    output_dir = os.path.dirname(os.path.abspath(output_file))
    assert output_dir != "", (
        f"Expected output_dir for output_file='{output_file}'  to be non-'' but found '{output_dir}'"
    )

    os.makedirs(output_dir, exist_ok=True)

    if adapter == "alora":
        # Tokenize the invocation string for PEFT 0.18.0 native aLoRA
        invocation_token_ids = tokenizer.encode(
            invocation_prompt, add_special_tokens=False
        )

        peft_config = LoraConfig(
            r=32,
            lora_alpha=32,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["q_proj", "k_proj", "v_proj"],
            alora_invocation_tokens=invocation_token_ids,  # Enable aLoRA
        )
        model = get_peft_model(model_base, peft_config)

        sft_args = SFTConfig(
            output_dir=output_dir,
            dataset_kwargs={"add_special_tokens": False},
            num_train_epochs=epochs,
            learning_rate=learning_rate,
            max_seq_length=max_length,
            per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=grad_accum,
            fp16=use_fp16,
        )

        trainer = SafeSaveTrainer(
            model,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
            args=sft_args,
            formatting_func=formatting_prompts_func,
            data_collator=collator,
            callbacks=[SaveBestModelCallback()],
        )
        trainer.train()
        model.save_pretrained(output_file)

    else:
        peft_config = LoraConfig(
            r=6,
            lora_alpha=32,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["q_proj", "k_proj", "v_proj"],
        )
        model = get_peft_model(model_base, peft_config)

        sft_args = SFTConfig(
            output_dir=output_dir,
            dataset_kwargs={"add_special_tokens": False},
            num_train_epochs=epochs,
            learning_rate=learning_rate,
            max_seq_length=max_length,
            per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=grad_accum,
            fp16=use_fp16,
        )

        trainer = SafeSaveTrainer(
            model,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
            args=sft_args,
            formatting_func=formatting_prompts_func,
            data_collator=collator,
        )
        trainer.train()
        model.save_pretrained(output_file, safe_serialization=True)

    # remove useless README file generated by trainer.
    # NOTE: this might delete an existing readme if the output_dir already existed.
    # But we don't do any git commits yet, so that's probably okay.
    annoying_readme = os.path.join(output_dir, "README.md")
    if os.path.exists(annoying_readme):
        print(
            f"WArning: assuming {annoying_readme} was a useless peft-generated README. Deleting."
        )
        os.remove(annoying_readme)
