from __future__ import annotations

import argparse
import json
from pathlib import Path

import yaml

from common import ROOT, localize_workspace_path

DEFAULT_CONFIG = ROOT / "plugins" / "1c" / "training" / "configs" / "qwen3-coder-30b-a3b-lora.yaml"


def fail_missing_dependencies(exc: Exception) -> None:
    raise SystemExit(
        "Missing training dependencies. Install `requirements-training.txt` "
        "or run inside the GPU training container. Original error: "
        f"{type(exc).__name__}: {exc}"
    )


def load_config(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as handle:
        data = yaml.safe_load(handle)
    if not isinstance(data, dict):
        raise ValueError(f"{path} must contain a YAML mapping")
    return data


def load_messages(path: Path) -> list[dict]:
    rows = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def build_training_texts(rows: list[dict], tokenizer) -> list[str]:
    texts = []
    for row in rows:
        messages = row["messages"]
        if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
        else:
            text = "\n".join(f"{msg['role']}: {msg['content']}" for msg in messages)
        texts.append(text)
    return texts


def main() -> int:
    parser = argparse.ArgumentParser(description="Train the 1C LoRA adapter.")
    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    config = load_config(args.config)
    base_model_path = str(localize_workspace_path(config["base_model_path"]))
    dataset_path = localize_workspace_path(config["dataset_path"])
    output_dir = str(localize_workspace_path(config["output_dir"]))

    rows = load_messages(dataset_path)
    if args.dry_run:
        print(f"Config: {args.config}")
        print(f"Base model: {base_model_path}")
        print(f"Dataset: {dataset_path} ({len(rows)} record(s))")
        print(f"Output: {output_dir}")
        return 0

    try:
        import torch
        from datasets import Dataset
        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
        from transformers import (
            AutoModelForCausalLM,
            AutoTokenizer,
            BitsAndBytesConfig,
            DataCollatorForLanguageModeling,
            Trainer,
            TrainingArguments,
        )
    except Exception as exc:  # pragma: no cover
        fail_missing_dependencies(exc)

    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for this training job.")

    tokenizer = AutoTokenizer.from_pretrained(base_model_path, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    qconf = config.get("quantization") or {}
    quantization_config = None
    if qconf.get("load_in_4bit"):
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type=qconf.get("bnb_4bit_quant_type", "nf4"),
            bnb_4bit_use_double_quant=bool(qconf.get("bnb_4bit_use_double_quant", True)),
            bnb_4bit_compute_dtype=torch.bfloat16,
        )

    model = AutoModelForCausalLM.from_pretrained(
        base_model_path,
        trust_remote_code=True,
        device_map="auto",
        quantization_config=quantization_config,
    )
    if quantization_config is not None:
        model = prepare_model_for_kbit_training(model)

    lora = config["lora"]
    peft_config = LoraConfig(
        r=int(lora["r"]),
        lora_alpha=int(lora["lora_alpha"]),
        lora_dropout=float(lora["lora_dropout"]),
        bias=lora.get("bias", "none"),
        task_type="CAUSAL_LM",
        target_modules=lora["target_modules"],
    )
    model = get_peft_model(model, peft_config)

    texts = build_training_texts(rows, tokenizer)
    dataset = Dataset.from_dict({"text": texts})

    def tokenize(batch):
        return tokenizer(
            batch["text"],
            truncation=True,
            max_length=int(config["max_seq_length"]),
            padding=False,
        )

    tokenized = dataset.map(tokenize, batched=True, remove_columns=["text"])
    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
    train = config["train"]
    train_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=float(train["num_train_epochs"]),
        per_device_train_batch_size=int(train["per_device_train_batch_size"]),
        gradient_accumulation_steps=int(train["gradient_accumulation_steps"]),
        learning_rate=float(train["learning_rate"]),
        warmup_ratio=float(train["warmup_ratio"]),
        logging_steps=int(train["logging_steps"]),
        save_strategy=train.get("save_strategy", "epoch"),
        save_total_limit=int(train.get("save_total_limit", 1)),
        bf16=bool(train.get("bf16", False)),
        fp16=bool(train.get("fp16", False)),
        report_to=[],
    )
    trainer = Trainer(model=model, args=train_args, train_dataset=tokenized, data_collator=collator)
    trainer.train()
    trainer.save_model(output_dir)
    tokenizer.save_pretrained(output_dir)
    print(f"Saved 1C LoRA adapter to {output_dir}")
    return 0


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