"""Train a psyc QLoRA adapter on JSONL Trainline datasets using unsloth + Qwen3.5.

Run inside the psyc training container:
    docker run --gpus all -v $(pwd)/data:/data psyc-trainer \
        --dataset /data/datasets/ioc_extraction-v1.jsonl \
        --dataset /data/datasets/severity_classification-v1.jsonl \
        --output /data/adapters/psyc-v1

Defaults target a 24 GB consumer GPU (3090/4090) with Qwen3.5-4B-Instruct at
4-bit + LoRA r=16. For an A100-40/80 bump --base-model to 9B and raise
--batch-size + --max-seq-length.
"""

from __future__ import annotations

# unsloth must be imported BEFORE transformers per their setup notes.
from unsloth import FastLanguageModel  # noqa: I001

import argparse
import json
from pathlib import Path
from typing import Dict, List

from datasets import Dataset
from trl import SFTConfig, SFTTrainer


def load_examples(paths: List[Path]) -> List[Dict[str, str]]:
    out: List[Dict[str, str]] = []
    for p in paths:
        with p.open("r", encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                ex = json.loads(line)
                if not all(k in ex for k in ("instruction", "input", "output")):
                    continue
                out.append(ex)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset", action="append", required=True, help="JSONL path (repeatable)")
    parser.add_argument("--base-model", default="unsloth/Qwen3.5-4B")
    parser.add_argument("--output", default="/data/adapters/psyc-v1")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--max-seq-length", type=int, default=4096)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--grad-accum", type=int, default=4)
    parser.add_argument("--lora-r", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=16)
    parser.add_argument("--seed", type=int, default=3407)
    args = parser.parse_args()

    paths = [Path(p) for p in args.dataset]
    examples = load_examples(paths)
    if not examples:
        raise SystemExit("no examples loaded — check --dataset paths")
    print(f"[psyc-train] loaded {len(examples)} example(s) from {len(paths)} dataset(s)")

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.base_model,
        max_seq_length=args.max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r=args.lora_r,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_alpha=args.lora_alpha,
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=args.seed,
    )

    def format_one(ex: Dict[str, str]) -> Dict[str, str]:
        messages = [
            {"role": "user", "content": f"{ex['instruction']}\n\n{ex['input']}"},
            {"role": "assistant", "content": ex["output"]},
        ]
        return {"text": tokenizer.apply_chat_template(messages, tokenize=False)}

    dataset = Dataset.from_list([format_one(e) for e in examples]).shuffle(seed=args.seed)

    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        args=SFTConfig(
            dataset_text_field="text",
            max_seq_length=args.max_seq_length,
            per_device_train_batch_size=args.batch_size,
            gradient_accumulation_steps=args.grad_accum,
            warmup_steps=5,
            num_train_epochs=args.epochs,
            learning_rate=args.lr,
            bf16=True,
            optim="adamw_8bit",
            weight_decay=0.01,
            lr_scheduler_type="linear",
            seed=args.seed,
            output_dir=str(output_dir / "checkpoints"),
            save_strategy="epoch",
            logging_steps=10,
            report_to="none",
        ),
    )
    train_result = trainer.train()

    final_dir = output_dir / "final"
    final_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(final_dir))
    tokenizer.save_pretrained(str(final_dir))

    loss_history = [
        {"step": h["step"], "loss": h["loss"], "epoch": h.get("epoch")}
        for h in trainer.state.log_history
        if "loss" in h
    ]
    (output_dir / "training_meta.json").write_text(json.dumps({
        "base_model": args.base_model,
        "lora_r": args.lora_r,
        "lora_alpha": args.lora_alpha,
        "epochs": args.epochs,
        "lr": args.lr,
        "datasets": [str(p) for p in paths],
        "examples": len(examples),
        "seed": args.seed,
        "train_loss": train_result.training_loss,
        "loss_history": loss_history,
    }, indent=2))
    print(f"[psyc-train] adapter saved → {final_dir}")


if __name__ == "__main__":
    main()