#!/usr/bin/env python3
"""
Production QLoRA training for BookBuilder bookwriter v1.
Same recipe as smoke_qlora.py, scaled up via HF Trainer.
Default target: Qwen 2.5 7B, 1.82B-token corpus, 1 epoch, ctx 2048.
"""
import os, sys, json, argparse, time
from pathlib import Path

import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig,
    Trainer, TrainingArguments, DataCollatorForLanguageModeling,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--base_model", default="Qwen/Qwen2.5-7B")
    p.add_argument("--train_jsonl", required=True)
    p.add_argument("--eval_jsonl", default=None)
    p.add_argument("--output_dir", required=True)
    p.add_argument("--max_seq", type=int, default=2048)
    p.add_argument("--per_device_train_batch_size", type=int, default=8)
    p.add_argument("--per_device_eval_batch_size", type=int, default=8)
    p.add_argument("--grad_accum", type=int, default=4)
    p.add_argument("--lr", type=float, default=2e-4)
    p.add_argument("--num_epochs", type=float, default=1.0)
    p.add_argument("--warmup_ratio", type=float, default=0.03)
    p.add_argument("--lora_r", type=int, default=16)
    p.add_argument("--lora_alpha", type=int, default=32)
    p.add_argument("--lora_dropout", type=float, default=0.05)
    p.add_argument("--save_steps", type=int, default=500)
    p.add_argument("--eval_steps", type=int, default=500)
    p.add_argument("--logging_steps", type=int, default=20)
    p.add_argument("--save_total_limit", type=int, default=3)
    p.add_argument("--hub_repo_id", default=None,
                   help="If set, push checkpoints to this HF repo each save.")
    p.add_argument("--hub_token", default=None)
    p.add_argument("--max_train_samples", type=int, default=None,
                   help="Cap for sanity runs.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--gradient_checkpointing", action="store_true", default=False)
    p.add_argument("--no_eval", action="store_true")
    return p.parse_args()


def main():
    args = parse_args()
    print(f"=== BookBuilder QLoRA train ===", flush=True)
    print(f"base: {args.base_model}")
    print(f"train: {args.train_jsonl}")
    print(f"eval : {args.eval_jsonl}")
    print(f"out  : {args.output_dir}")
    print(f"ctx  : {args.max_seq}  bs/dev: {args.per_device_train_batch_size}  ga: {args.grad_accum}  lr: {args.lr}", flush=True)

    # ---------- tokenizer ----------
    tok = AutoTokenizer.from_pretrained(args.base_model, use_fast=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "right"

    # ---------- model: 4-bit NF4 ----------
    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    print("Loading 4-bit base...", flush=True)
    t0 = time.time()
    model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        quantization_config=bnb,
        torch_dtype=torch.bfloat16,
        attn_implementation="sdpa",
        device_map="cuda:0",
    )
    model.config._attn_implementation = "sdpa"
    model.config.use_cache = False
    print(f"  base loaded in {time.time()-t0:.1f}s | VRAM {torch.cuda.memory_allocated()/1e9:.2f} GB", flush=True)

    model = prepare_model_for_kbit_training(
        model, use_gradient_checkpointing=args.gradient_checkpointing,
    )
    lora = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
    )
    model = get_peft_model(model, lora)
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total     = sum(p.numel() for p in model.parameters())
    print(f"  trainable: {trainable/1e6:.2f}M / {total/1e6:.1f}M ({100*trainable/total:.3f}%)", flush=True)

    # ---------- dataset ----------
    data_files = {"train": args.train_jsonl}
    if args.eval_jsonl and not args.no_eval:
        data_files["eval"] = args.eval_jsonl
    ds = load_dataset("json", data_files=data_files)
    if args.max_train_samples:
        ds["train"] = ds["train"].shuffle(seed=args.seed).select(range(args.max_train_samples))
    print(f"  train rows: {len(ds['train']):,}", flush=True)
    if "eval" in ds:
        print(f"  eval rows : {len(ds['eval']):,}", flush=True)

    def tokenize_fn(batch):
        out = tok(
            batch["text"],
            truncation=True,
            max_length=args.max_seq,
            padding=False,
        )
        return out

    # Explicit cache paths — HF datasets' default fingerprint isn't stable across
    # Python processes (tokenizer object hash differs), so a plain .map() re-tokenizes
    # every fresh launch. Pin the cache to a deterministic path so subsequent runs hit it.
    cache_root = os.path.dirname(os.path.abspath(args.train_jsonl))
    tag = f"qwen-ctx{args.max_seq}"
    cache_files = {
        split: os.path.join(cache_root, f"tok_cache_{split}_{tag}.arrow")
        for split in ds.keys()
    }
    for split, path in cache_files.items():
        print(f"  tokenize cache[{split}] -> {path}", flush=True)
    tokenized = ds.map(
        tokenize_fn,
        batched=True,
        remove_columns=ds["train"].column_names,
        num_proc=8,
        desc="tokenize",
        cache_file_names=cache_files,
        load_from_cache_file=True,
    )

    collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False)

    # ---------- trainer ----------
    push_to_hub = bool(args.hub_repo_id)
    targs = TrainingArguments(
        output_dir=args.output_dir,
        overwrite_output_dir=False,
        seed=args.seed,
        num_train_epochs=args.num_epochs,
        per_device_train_batch_size=args.per_device_train_batch_size,
        per_device_eval_batch_size=args.per_device_eval_batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        warmup_ratio=args.warmup_ratio,
        lr_scheduler_type="cosine",
        bf16=True,
        fp16=False,
        tf32=True,
        gradient_checkpointing=args.gradient_checkpointing,
        logging_steps=args.logging_steps,
        save_steps=args.save_steps,
        eval_strategy="steps" if "eval" in tokenized else "no",
        eval_steps=args.eval_steps if "eval" in tokenized else None,
        save_total_limit=args.save_total_limit,
        report_to="tensorboard",
        push_to_hub=push_to_hub,
        hub_model_id=args.hub_repo_id if push_to_hub else None,
        hub_token=args.hub_token,
        hub_strategy="every_save" if push_to_hub else "end",
        hub_private_repo=False,
        dataloader_num_workers=4,
        dataloader_pin_memory=True,
        optim="paged_adamw_8bit",
        remove_unused_columns=False,
    )

    trainer_kwargs = dict(
        model=model,
        args=targs,
        train_dataset=tokenized["train"],
        eval_dataset=tokenized.get("eval"),
        data_collator=collator,
    )
    # transformers 4.46 uses `tokenizer=`; 5.x uses `processing_class=`
    import inspect
    if "processing_class" in inspect.signature(Trainer.__init__).parameters:
        trainer_kwargs["processing_class"] = tok
    else:
        trainer_kwargs["tokenizer"] = tok
    trainer = Trainer(**trainer_kwargs)

    print(f"=== begin training ===", flush=True)
    trainer.train()

    print(f"=== saving final adapter ===", flush=True)
    trainer.save_model(args.output_dir)
    tok.save_pretrained(args.output_dir)
    if push_to_hub:
        trainer.push_to_hub(commit_message="final adapter")

    # Sentinel
    Path(args.output_dir).joinpath(".train_complete").touch()
    print("DONE.", flush=True)


if __name__ == "__main__":
    main()
