"""LoRA fine-tuning of Qwen3.5 for runic recognition (image -> transliteration) or
translation (transliteration -> English/Swedish). Loss is computed on the answer tokens only.

Data (jsonl):
  mt : {"id", "src", "tgt", "lang"}
  ocr: {"id", "image", "target"}
"""
import argparse
import json
import math
import os
import random
import sys
import time

import torch
from torch.utils.data import DataLoader, Dataset

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from ocr_noise import corrupt  # noqa: E402
from qwen_common import encode_kwargs, end_text, load_image, load_model, prompt_text  # noqa: E402


def read_jsonl(paths, limit=None):
    rows = []
    for p in paths.split(","):
        with open(p) as f:
            rows.extend(json.loads(l) for l in f if l.strip())
    if limit:
        rows = rows[:limit]
    return rows


class Examples(Dataset):
    def __init__(self, rows, processor, task, image_root, noise_p=0.0, noise_max=0.5, seed=0, tags=False):
        self.rows, self.processor, self.task, self.tags = rows, processor, task, tags
        self.image_root, self.noise_p, self.noise_max = image_root, noise_p, noise_max
        self.rng = random.Random(seed)
        self.end = end_text(processor.tokenizer)
        self.kw = encode_kwargs(processor)

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

    def __getitem__(self, i):
        r = self.rows[i]
        if self.task == "ocr":
            img = load_image(os.path.join(self.image_root, r["image"]))
            prompt = prompt_text(self.processor, "ocr")
            target = r["target"]
            full = self.processor(text=[prompt + target + self.end], images=[img], return_tensors="pt")
            n_prompt = self.processor(text=[prompt], images=[img], return_tensors="pt")["input_ids"].shape[1]
        else:
            src = r["src"]
            rng = random.Random(int.from_bytes(os.urandom(8), "little"))
            if self.noise_p > 0 and rng.random() < self.noise_p:
                src = corrupt(src, rng.uniform(0.05, self.noise_max), rng)
            prompt = prompt_text(self.processor, "mt", src, r["lang"], r["rl"] if self.tags else None)
            target = r["tgt"]
            full = self.processor(text=[prompt + target + self.end], return_tensors="pt", **self.kw)
            n_prompt = self.processor(text=[prompt], return_tensors="pt", **self.kw)["input_ids"].shape[1]
        item = {k: v for k, v in full.items()}
        labels = item["input_ids"].clone()
        labels[:, :n_prompt] = -100
        item["labels"] = labels
        return item


def collate(batch, pad_id, max_len):
    out = {}
    seq_keys = [k for k, v in batch[0].items() if v.dim() == 2 and v.shape[0] == 1 and k != "image_grid_thw"]
    L = min(max(b["input_ids"].shape[1] for b in batch), max_len)
    for k in seq_keys:
        pad = pad_id if k == "input_ids" else (-100 if k == "labels" else 0)
        rows = []
        for b in batch:
            t = b[k][0, :L]
            rows.append(torch.cat([t, torch.full((L - t.shape[0],), pad, dtype=t.dtype)]))
        out[k] = torch.stack(rows)
    for k in batch[0]:
        if k not in seq_keys:
            out[k] = torch.cat([b[k] for b in batch], dim=0)
    return out


@torch.no_grad()
def val_loss(model, loader, device, max_batches=50):
    model.eval()
    tot, n = 0.0, 0
    for bi, batch in enumerate(loader):
        if bi >= max_batches:
            break
        batch = {k: v.to(device) for k, v in batch.items()}
        with torch.autocast("cuda", dtype=torch.bfloat16):
            tot += model(**batch).loss.item()
        n += 1
    model.train()
    return tot / max(n, 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--task", choices=["mt", "ocr"], required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--train_file", required=True)
    ap.add_argument("--val_file", required=True)
    ap.add_argument("--image_root", default="")
    ap.add_argument("--out", required=True)
    ap.add_argument("--epochs", type=float, default=3)
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--bs", type=int, default=8)
    ap.add_argument("--grad_accum", type=int, default=2)
    ap.add_argument("--rank", type=int, default=16)
    ap.add_argument("--alpha", type=int, default=32)
    ap.add_argument("--max_len", type=int, default=1536)
    ap.add_argument("--max_steps", type=int, default=0)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--noise_p", type=float, default=0.0)
    ap.add_argument("--seed", type=int, default=13)
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--init_adapter", default="")
    ap.add_argument("--tags", action="store_true", help="mt: prepend <runic-language> <2target> tags (field rl)")
    ap.add_argument("--no_last", action="store_true", help="do not save last_adapter (disk)")
    args = ap.parse_args()
    print("ARGS", vars(args), flush=True)
    torch.manual_seed(args.seed)
    random.seed(args.seed)
    os.makedirs(args.out, exist_ok=True)

    processor, model = load_model(args.model, adapter=args.init_adapter or None, trainable=bool(args.init_adapter))
    tok = processor.tokenizer
    if not args.init_adapter:
        from peft import LoraConfig, get_peft_model

        cfg = LoraConfig(r=args.rank, lora_alpha=args.alpha, lora_dropout=0.05, target_modules="all-linear",
                         task_type="CAUSAL_LM")
        model = get_peft_model(model, cfg)
    model.print_trainable_parameters()
    print("base dtype", next(p for n, p in model.named_parameters() if "lora" not in n).dtype,
          "lora dtype", next(p for n, p in model.named_parameters() if "lora" in n).dtype, flush=True)
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    model.enable_input_require_grads()
    device = "cuda"
    model.to(device)
    model.train()

    tr_rows = read_jsonl(args.train_file, args.limit or None)
    va_rows = read_jsonl(args.val_file, 400 if not args.limit else min(args.limit, 40))
    random.Random(args.seed).shuffle(tr_rows)
    tr = Examples(tr_rows, processor, args.task, args.image_root, args.noise_p, seed=args.seed, tags=args.tags)
    va = Examples(va_rows, processor, args.task, args.image_root, 0.0, seed=args.seed, tags=args.tags)
    col = lambda b: collate(b, tok.pad_token_id, args.max_len)  # noqa: E731
    tl = DataLoader(tr, batch_size=args.bs, shuffle=True, num_workers=args.workers, collate_fn=col,
                    persistent_workers=args.workers > 0)
    vl = DataLoader(va, batch_size=args.bs, shuffle=False, num_workers=args.workers, collate_fn=col)

    params = [p for p in model.parameters() if p.requires_grad]
    opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.0)
    steps_per_epoch = math.ceil(len(tl) / args.grad_accum)
    total = args.max_steps or int(steps_per_epoch * args.epochs)
    warm = max(1, int(0.03 * total))
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min(1.0, (s + 1) / warm) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / max(1, total)))))
    print(f"train {len(tr)} val {len(va)} steps/epoch {steps_per_epoch} total {total}", flush=True)

    step, best, t0 = 0, float("inf"), time.time()
    log = open(os.path.join(args.out, "train_log.jsonl"), "a")
    done = False
    epoch = 0
    while not done:
        for bi, batch in enumerate(tl):
            batch = {k: v.to(device) for k, v in batch.items()}
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss = model(**batch).loss / args.grad_accum
            loss.backward()
            if (bi + 1) % args.grad_accum:
                continue
            torch.nn.utils.clip_grad_norm_(params, 1.0)
            opt.step()
            sched.step()
            opt.zero_grad(set_to_none=True)
            step += 1
            if step % 20 == 0 or step == 1:
                rec = dict(step=step, loss=loss.item() * args.grad_accum, lr=sched.get_last_lr()[0],
                           elapsed=round(time.time() - t0))
                print(json.dumps(rec), flush=True)
                log.write(json.dumps(rec) + "\n")
                log.flush()
            eval_every = max(1, steps_per_epoch // 2) if not args.max_steps else args.max_steps
            if step % eval_every == 0 or step >= total:
                vloss = val_loss(model, vl, device)
                rec = dict(step=step, epoch=round(step / steps_per_epoch, 2), val_loss=vloss)
                print(json.dumps(rec), flush=True)
                log.write(json.dumps(rec) + "\n")
                log.flush()
                if vloss < best:
                    best = vloss
                    model.save_pretrained(os.path.join(args.out, "best_adapter"))
                    json.dump(rec, open(os.path.join(args.out, "best.json"), "w"))
            if step >= total:
                done = True
                break
        epoch += 1
    if not args.no_last:
        model.save_pretrained(os.path.join(args.out, "last_adapter"))
    json.dump(dict(total_steps=step, best_val_loss=best, minutes=round((time.time() - t0) / 60, 1)),
              open(os.path.join(args.out, "DONE.json"), "w"))
    print("DONE", step, best, flush=True)


if __name__ == "__main__":
    main()