#!/usr/bin/env python3
"""LoRA SFT for Qwen3-Coder on the code -> Mermaid flowchart dataset.

The dataset (data/train.jsonl, data/val.jsonl) is chat "messages" JSONL with a
system / user / assistant turn per line. We train only on the assistant turn
(the prompt tokens are masked to -100) so the model learns to *produce* the
<thinking> + Mermaid graph + <linemap>, not to echo the prompt.

Quick checks (no 30B download needed), using the project venv:
    .venv/bin/python finetune.py --dry-run            # inspect token stats
Full run (do this on a CUDA GPU box; 30B needs 4-bit + a big card):
    .venv/bin/python finetune.py --4bit --epochs 3 --output-dir out/qwen-mermaid

Note on the target model: "Qwen3-Coder-30B-A3B-Instruct-UD-Q3_K_XL" is an
Unsloth Dynamic GGUF *inference* quant. The public llama.cpp export path can
produce standard GGUF quants such as Q3_K_XL, but not Unsloth's model-specific
UD-* dynamic quant recipes. See README.md.
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
from typing import Any, Callable, cast

DEFAULT_MODEL = "Qwen/Qwen3-Coder-30B-A3B-Instruct"
HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_LLAMA_CPP_DIR = os.path.join(HERE, ".venv", "llama.cpp")
DEFAULT_LLAMA_CPP_REPO = "https://github.com/ggml-org/llama.cpp.git"


# --------------------------------------------------------------------------- #
# Data
# --------------------------------------------------------------------------- #


def read_jsonl(path):
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh if line.strip()]


def encode_example(tokenizer, messages, max_seq_len):
    """Return (input_ids, labels) with the prompt masked to -100.

    The full conversation is encoded once; the prompt is the same conversation
    minus the final assistant turn, with a generation prompt appended. Qwen's
    chat template is prefix-consistent, so the first len(prompt) tokens of the
    full encoding are exactly the prompt — we mask those.
    """
    full = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)
    prompt = tokenizer.apply_chat_template(messages[:-1], tokenize=True, add_generation_prompt=True)
    full = full[:max_seq_len]
    n_prompt = min(len(prompt), len(full))
    labels = [-100] * n_prompt + full[n_prompt:]
    return full, labels


class ChatDataset:
    """Lazy torch dataset over messages JSONL with prompt-masked labels."""

    def __init__(self, path, tokenizer, max_seq_len):
        self.records = read_jsonl(path)
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len

    def __len__(self):
        return len(self.records)

    def __getitem__(self, i):
        ids, labels = encode_example(self.tokenizer, self.records[i]["messages"], self.max_seq_len)
        return {"input_ids": ids, "labels": labels, "attention_mask": [1] * len(ids)}


class Collator:
    """Right-pad input_ids/attention_mask/labels to the batch max."""

    def __init__(self, pad_id):
        self.pad_id = pad_id

    def __call__(self, batch):
        import torch
        width = max(len(b["input_ids"]) for b in batch)
        out = {"input_ids": [], "attention_mask": [], "labels": []}
        for b in batch:
            pad = width - len(b["input_ids"])
            out["input_ids"].append(b["input_ids"] + [self.pad_id] * pad)
            out["attention_mask"].append(b["attention_mask"] + [0] * pad)
            out["labels"].append(b["labels"] + [-100] * pad)
        return {k: torch.tensor(v, dtype=torch.long) for k, v in out.items()}


# --------------------------------------------------------------------------- #
# Dry run: validate the data pipeline without loading the 30B model
# --------------------------------------------------------------------------- #


def dry_run(args):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(args.tokenizer or args.model, trust_remote_code=True)
    ds = ChatDataset(args.train, tok, args.max_seq_len)
    lens, masked, supervised = [], 0, 0
    for i in range(min(len(ds), 400)):
        ex = ds[i]
        lens.append(len(ex["input_ids"]))
        masked += sum(1 for x in ex["labels"] if x == -100)
        supervised += sum(1 for x in ex["labels"] if x != -100)
    lens.sort()
    print(f"examples            : {len(ds)} (sampled {len(lens)})")
    print(f"seq length p50/p95/max: {lens[len(lens)//2]} / {lens[int(len(lens)*0.95)]} / {lens[-1]}")
    print(f"supervised tokens   : {supervised}  masked(prompt) tokens: {masked}")
    print(f"max_seq_len {args.max_seq_len} -> "
          f"{sum(1 for l in lens if l >= args.max_seq_len)} sampled examples hit the cap")
    print("\nsample decoded labels (assistant target the model is trained to emit):")
    ex = ds[0]
    sup = [t for t, l in zip(ex["input_ids"], ex["labels"]) if l != -100]
    print(tok.decode(sup)[:600])


# --------------------------------------------------------------------------- #
# Merge / GGUF export
# --------------------------------------------------------------------------- #


def _run(cmd, cwd=None):
    print("+ " + " ".join(str(x) for x in cmd), flush=True)
    subprocess.run(cmd, cwd=cwd, check=True)


def _find_llama_cpp_file(root, names):
    for name in names:
        path = os.path.join(root, name)
        if os.path.exists(path):
            return path
    return None


def _which(name):
    venv_bin = os.path.dirname(sys.executable)
    candidate = os.path.join(venv_bin, name)
    if os.path.exists(candidate):
        return candidate
    return shutil.which(name)


def cleanup_path(path, label):
    if not path:
        return
    path = os.path.abspath(os.path.expanduser(path))
    if not os.path.exists(path):
        return
    print(f"removing {label} to free disk -> {path}", flush=True)
    shutil.rmtree(path, ignore_errors=True)


def ensure_llama_cpp(args):
    """Return a llama.cpp checkout with convert + quantize available.

    There is no pip wheel that provides the full llama.cpp CLI toolchain we need
    for GGUF export. If the checkout/binary is missing, build it inside the venv.
    """
    llama_dir = os.path.abspath(args.llama_cpp_dir or DEFAULT_LLAMA_CPP_DIR)
    convert = _find_llama_cpp_file(llama_dir, ["convert_hf_to_gguf.py"])
    quantize = _find_llama_cpp_file(
        llama_dir,
        [
            "build/bin/llama-quantize",
            "build/bin/quantize",
            "llama-quantize",
            "quantize",
        ],
    )
    if convert and quantize:
        return llama_dir

    if args.no_auto_build_llama_cpp:
        raise SystemExit(
            f"llama.cpp is incomplete under {llama_dir}. Re-run without "
            "--no-auto-build-llama-cpp or build llama.cpp manually."
        )

    git = _which("git")
    cmake = _which("cmake")
    if git is None:
        raise SystemExit("git is required to auto-clone llama.cpp")
    if cmake is None:
        raise SystemExit("cmake is required to auto-build llama.cpp; install requirements first")

    if not os.path.exists(llama_dir):
        os.makedirs(os.path.dirname(llama_dir), exist_ok=True)
        _run([git, "clone", args.llama_cpp_repo, llama_dir])
    elif not os.path.exists(os.path.join(llama_dir, "CMakeLists.txt")):
        raise SystemExit(
            f"{llama_dir} exists but does not look like a llama.cpp checkout"
        )

    build_dir = os.path.join(llama_dir, "build")
    _run([
        cmake,
        "-S", llama_dir,
        "-B", build_dir,
        "-DGGML_CUDA=ON",
        "-DCMAKE_BUILD_TYPE=Release",
    ])
    _run([
        cmake,
        "--build", build_dir,
        "--config", "Release",
        "-j",
        "--target", "llama-quantize",
    ])

    convert = _find_llama_cpp_file(llama_dir, ["convert_hf_to_gguf.py"])
    quantize = _find_llama_cpp_file(llama_dir, ["build/bin/llama-quantize", "build/bin/quantize"])
    if not convert or not quantize:
        raise SystemExit(f"failed to build llama.cpp export tools under {llama_dir}")
    return llama_dir


def merge_lora(args):
    """Merge the saved adapter into a normal HF checkpoint for conversion."""
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    dtype = torch.bfloat16 if args.merge_dtype == "bfloat16" else torch.float16
    merge_dir = os.path.abspath(args.merge_dir)
    os.makedirs(merge_dir, exist_ok=True)

    tok = AutoTokenizer.from_pretrained(args.output_dir, trust_remote_code=True)
    tok.save_pretrained(merge_dir)

    device_map = {"": "cpu"} if args.merge_device_map == "cpu" else args.merge_device_map
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        trust_remote_code=True,
        torch_dtype=dtype,
        device_map=device_map,
        low_cpu_mem_usage=True,
    )
    model = PeftModel.from_pretrained(model, args.output_dir)
    merge_and_unload = cast(Callable[[], Any], getattr(model, "merge_and_unload"))
    model = merge_and_unload()
    model.save_pretrained(
        merge_dir,
        safe_serialization=True,
        max_shard_size=args.max_shard_size,
    )
    print(f"saved merged HF model -> {merge_dir}")
    return merge_dir


def export_gguf(args, merged_dir):
    if args.gguf_quant.upper().startswith("UD-"):
        raise SystemExit(
            "Unsloth Dynamic GGUF quant names such as UD-Q3_K_XL are not "
            "llama.cpp quantization types. Export a standard quant such as "
            "Q3_K_XL here, or use an official Unsloth Dynamic GGUF release."
        )

    llama_dir = ensure_llama_cpp(args)
    convert = _find_llama_cpp_file(llama_dir, ["convert_hf_to_gguf.py"])
    if not convert:
        raise SystemExit(f"could not find convert_hf_to_gguf.py under {llama_dir}")

    quantize = _find_llama_cpp_file(
        llama_dir,
        [
            "build/bin/llama-quantize",
            "build/bin/quantize",
            "llama-quantize",
            "quantize",
        ],
    )
    if not quantize:
        raise SystemExit(
            "could not find a llama.cpp quantize binary. Build llama.cpp first, "
            "for example: cmake -B build -DGGML_CUDA=ON && cmake --build build -j"
        )

    out_path = os.path.abspath(args.gguf_out)
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    f16_path = os.path.splitext(out_path)[0] + ".F16.gguf"

    _run([sys.executable or "python", convert, merged_dir, "--outfile", f16_path, "--outtype", args.gguf_outtype])
    _run([quantize, f16_path, out_path, args.gguf_quant])

    if args.keep_f16_gguf:
        print(f"saved intermediate GGUF -> {f16_path}")
    else:
        try:
            os.remove(f16_path)
        except OSError:
            pass

    print(f"saved quantized GGUF -> {out_path}")
    return out_path


# --------------------------------------------------------------------------- #
# Training
# --------------------------------------------------------------------------- #


def train(args):
    import torch
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from transformers import (AutoModelForCausalLM, AutoTokenizer, Trainer,
                              TrainingArguments)

    tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token

    model_kwargs: dict[str, Any] = {"trust_remote_code": True}
    if args.bf16:
        model_kwargs["torch_dtype"] = torch.bfloat16
    if args.fourbit:
        from transformers import BitsAndBytesConfig
        model_kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
        model_kwargs["device_map"] = "auto"

    model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)
    model.config.use_cache = False
    if args.fourbit:
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
    else:
        model.gradient_checkpointing_enable()

    lora = LoraConfig(
        r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=0.05, bias="none",
        task_type="CAUSAL_LM",
        # attention + MLP projections; works for the Qwen3-MoE A3B blocks too.
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
    )
    model = get_peft_model(model, lora)
    model.print_trainable_parameters()

    train_ds = ChatDataset(args.train, tok, args.max_seq_len)
    eval_ds = ChatDataset(args.val, tok, args.max_seq_len) if os.path.exists(args.val) else None

    targs = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        logging_steps=10,
        save_strategy="epoch",
        eval_strategy="epoch" if eval_ds else "no",
        bf16=args.bf16,
        gradient_checkpointing=not args.fourbit,
        report_to=[],
    )
    trainer = Trainer(model=model, args=targs, train_dataset=train_ds,
                      eval_dataset=eval_ds, data_collator=Collator(tok.pad_token_id))
    trainer.train()
    trainer.save_model(args.output_dir)
    tok.save_pretrained(args.output_dir)
    print(f"saved LoRA adapter -> {args.output_dir}")

    if args.merge or args.export_gguf:
        del trainer, model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        merged_dir = merge_lora(args)
        if args.export_gguf:
            if args.delete_hf_cache_before_gguf:
                cleanup_path(os.environ.get("HF_HOME"), "HF_HOME cache")
                cleanup_path(os.environ.get("TRANSFORMERS_CACHE"), "TRANSFORMERS_CACHE")
                cleanup_path(os.environ.get("HF_HUB_CACHE"), "HF_HUB_CACHE")
            export_gguf(args, merged_dir)
            if args.delete_merged_after_gguf:
                cleanup_path(merged_dir, "merged HF checkpoint")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--tokenizer", default=None, help="override tokenizer (dry-run convenience)")
    ap.add_argument("--train", default=os.path.join(HERE, "data/train.jsonl"))
    ap.add_argument("--val", default=os.path.join(HERE, "data/val.jsonl"))
    ap.add_argument("--output-dir", default=os.path.join(HERE, "out/qwen-mermaid-lora"))
    ap.add_argument("--epochs", type=float, default=3.0)
    ap.add_argument("--batch-size", type=int, default=1)
    ap.add_argument("--grad-accum", type=int, default=16)
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--max-seq-len", type=int, default=2048)
    ap.add_argument("--lora-r", type=int, default=16)
    ap.add_argument("--lora-alpha", type=int, default=32)
    ap.add_argument("--4bit", dest="fourbit", action="store_true", help="QLoRA (needs CUDA + bitsandbytes)")
    ap.add_argument("--bf16", action="store_true", default=True)
    ap.add_argument("--dry-run", action="store_true", help="inspect tokenized data, skip model load")
    ap.add_argument("--merge", action="store_true", help="merge the LoRA adapter into a HF checkpoint after training")
    ap.add_argument("--merge-dir", default=os.path.join(HERE, "out/qwen-mermaid-merged"))
    ap.add_argument("--merge-dtype", choices=["float16", "bfloat16"], default="float16")
    ap.add_argument("--merge-device-map", default="auto",
                    help='device_map for merge reload; use "cpu" if GPU RAM is tight')
    ap.add_argument("--max-shard-size", default="4GB")
    ap.add_argument("--export-gguf", action="store_true",
                    help="merge and export a quantized GGUF with llama.cpp")
    ap.add_argument("--llama-cpp-dir", default=os.environ.get("LLAMA_CPP_DIR"),
                    help="llama.cpp checkout; defaults to .venv/llama.cpp")
    ap.add_argument("--llama-cpp-repo", default=DEFAULT_LLAMA_CPP_REPO)
    ap.add_argument("--no-auto-build-llama-cpp", action="store_true")
    ap.add_argument("--gguf-out", default=os.path.join(HERE, "out/qwen3-coder-codeflow-Q3_K_XL.gguf"))
    ap.add_argument("--gguf-outtype", choices=["f32", "f16", "bf16", "q8_0"], default="f16")
    ap.add_argument("--gguf-quant", default="Q3_K_XL")
    ap.add_argument("--keep-f16-gguf", action="store_true")
    ap.add_argument("--delete-hf-cache-before-gguf", action="store_true",
                    help="delete local HF caches after merge and before GGUF conversion to reduce disk peak")
    ap.add_argument("--delete-merged-after-gguf", action="store_true",
                    help="delete the merged HF checkpoint after the final GGUF is written")
    args = ap.parse_args()

    if args.export_gguf:
        args.merge = True

    if args.dry_run:
        dry_run(args)
    else:
        train(args)


if __name__ == "__main__":
    main()