"""Stage 1 continual pre-training (PRIVATE to Training_Module).

CPU tiers (2 and 3): Layer-by-Layer ELEVATOR unfreezing - only 2 transformer layers are trainable
                     at any moment (promoted to fp32), all the others stay frozen. The window rides
                     from layer 0 up to layer N-1.  In tier 3 the window layers are pinned in the ring
                     loader and persisted to disk when the elevator leaves them.
GPU tier (1)       : rsLoRA on the 4-bit NF4 base model.

Data: ./cleaned_data/pretrain_clean.jsonl  ->  ~75% repo text chunks + ~25% rational/math/identity chunks
"""
from __future__ import annotations

import contextlib
import gc
import json
import math
import random
import time
from pathlib import Path
from typing import Callable, Dict, Iterator, List, Optional

import torch

LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]


# ------------------------------------------------------------------- helpers
def model_device(model) -> torch.device:
    try:
        return next(model.parameters()).device
    except StopIteration:
        return torch.device("cpu")


def make_amp(plan: Dict):
    """Return (autocast_factory, grad_scaler_or_None) for the plan."""
    dtype = plan["dtype"]
    if plan["tier"] == 1 and torch.cuda.is_available():
        dt = getattr(torch, dtype)
        scaler = None
        if dt == torch.float16:
            try:
                scaler = torch.amp.GradScaler("cuda")
            except Exception:
                scaler = torch.cuda.amp.GradScaler()
        return (lambda: torch.autocast("cuda", dtype=dt)), scaler
    if dtype == "bfloat16":
        return (lambda: torch.autocast("cpu", dtype=torch.bfloat16)), None
    return (lambda: contextlib.nullcontext()), None


def find_layers(model):
    base = model.get_base_model() if hasattr(model, "get_base_model") else model
    return base.model.layers


def train_loop(model, params: List[torch.nn.Parameter], batches: Iterator[Dict], loss_fn: Callable,
               steps: int, lr: float, accum: int, amp, scaler, tag: str, log_every: int = 1) -> float:
    """Generic AdamW loop with warmup + cosine decay, grad clipping and ETA logging."""
    opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.0, betas=(0.9, 0.95))
    warm = max(1, int(0.1 * steps))

    def lr_at(step: int) -> float:
        if step < warm:
            return (step + 1) / warm
        progress = (step - warm) / max(1, steps - warm)
        return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress))

    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_at)
    model.train()
    started = time.time()
    history: List[float] = []
    for step in range(steps):
        total = 0.0
        for _ in range(accum):
            batch = next(batches)
            with amp():
                loss = loss_fn(model, batch) / accum
            if scaler is not None:
                scaler.scale(loss).backward()
            else:
                loss.backward()
            total += float(loss.detach())
        if scaler is not None:
            scaler.unscale_(opt)
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        if scaler is not None:
            scaler.step(opt)
            scaler.update()
        else:
            opt.step()
        sched.step()
        opt.zero_grad(set_to_none=True)
        history.append(total)
        if (step + 1) % log_every == 0 or step == steps - 1:
            elapsed = time.time() - started
            eta = elapsed / (step + 1) * (steps - step - 1)
            print(f"   [{tag}] step {step + 1:>4}/{steps} loss {total:.4f} lr {opt.param_groups[0]['lr']:.2e} "
                  f"elapsed {elapsed:5.0f}s eta {eta:5.0f}s")
    del opt, sched
    gc.collect()
    tail = history[len(history) // 2:] or history
    return sum(tail) / max(len(tail), 1)


# -------------------------------------------------------------------- data
def build_chunks(tokenizer, path: Path, seq_len: int, max_chunks: int, conv_ratio: float,
                 seed: int = 1337, max_chunks_per_doc: int = 8) -> List[List[int]]:
    """Tokenise pretrain_clean.jsonl into fixed windows; chat.jsonl chunks make up ~conv_ratio."""
    eos = tokenizer.convert_tokens_to_ids("<|endoftext|>")
    if eos is None or eos == tokenizer.unk_token_id:
        eos = tokenizer.eos_token_id
    n_chat_target = int(round(max_chunks * conv_ratio))
    n_repo_target = max_chunks - n_chat_target
    repo_chunks: List[List[int]] = []
    chat_chunks: List[List[int]] = []
    buffers = {"repo": [], "chat": []}
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            rec = json.loads(line)
            kind = "chat" if rec.get("source") == "chat.jsonl" else "repo"
            bucket = chat_chunks if kind == "chat" else repo_chunks
            target = n_chat_target if kind == "chat" else n_repo_target
            if len(bucket) >= max(target, 1) * (3 if kind == "chat" else 1):
                if len(repo_chunks) >= n_repo_target and len(chat_chunks) >= n_chat_target:
                    break
                continue
            ids = tokenizer(rec["text"], add_special_tokens=False, truncation=False)["input_ids"]
            ids = ids[: seq_len * max_chunks_per_doc] + [eos]
            buf = buffers[kind]
            buf.extend(ids)
            while len(buf) >= seq_len:
                bucket.append(buf[:seq_len])
                del buf[:seq_len]
    rng = random.Random(seed)
    rng.shuffle(repo_chunks)
    rng.shuffle(chat_chunks)
    repo_part = repo_chunks[:n_repo_target]
    if not chat_chunks:
        chat_part: List[List[int]] = []
    elif not repo_part:
        chat_part = [chat_chunks[i % len(chat_chunks)] for i in range(max_chunks)]
    else:
        chat_part = [chat_chunks[i % len(chat_chunks)] for i in range(n_chat_target)]
    mixed = repo_part + chat_part
    rng.shuffle(mixed)
    share = 100.0 * len(chat_part) / max(len(mixed), 1)
    print(f"   [pretrain-data] {len(mixed)} chunks of {seq_len} tokens "
          f"(repo={len(repo_part)}, chat={len(chat_part)} -> {share:.0f}% conversation share)")
    return mixed


def batch_stream(chunks: List[List[int]], batch_size: int, device, seed: int = 7) -> Iterator[Dict]:
    rng = random.Random(seed)
    order: List[int] = []
    while True:
        if len(order) < batch_size:
            fresh = list(range(len(chunks)))
            rng.shuffle(fresh)
            order.extend(fresh)
        take, order = order[:batch_size], order[batch_size:]
        ids = torch.tensor([chunks[i] for i in take], dtype=torch.long, device=device)
        yield {"input_ids": ids, "attention_mask": torch.ones_like(ids), "labels": ids.clone()}


def lm_loss(model, batch):
    out = model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["labels"])
    return out.loss


# -------------------------------------------------------------------- LoRA
def attach_lora(model, plan: Dict, rank: int, rslora: bool, resized: bool, adapter_dir: Optional[str] = None,
                save_resized_modules: bool = True):
    from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
    if plan.get("quant") == "nf4":
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True,
                                                gradient_checkpointing_kwargs={"use_reentrant": False})
    if adapter_dir:
        print(f"   [lora] continuing adapter from {adapter_dir}")
        return PeftModel.from_pretrained(model, adapter_dir, is_trainable=True)
    cfg = LoraConfig(r=rank, lora_alpha=rank if rslora else 2 * rank, lora_dropout=0.05,
                     target_modules=LORA_TARGETS, bias="none", task_type="CAUSAL_LM", use_rslora=rslora,
                     modules_to_save=["embed_tokens", "lm_head"] if (resized and save_resized_modules) else None)
    model = get_peft_model(model, cfg)
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"   [lora] {'rsLoRA' if rslora else 'LoRA'} r={rank}: {trainable / 1e6:.2f}M trainable params")
    return model


def save_stage1_checkpoint(model, tokenizer, out_dir: Path, cfg: Dict, kind: str) -> Path:
    """Write Stage 1 to a sibling temp directory and publish it without deleting an existing checkpoint."""
    out_dir = Path(out_dir)
    tmp_dir = out_dir.with_name(out_dir.name + ".tmp")
    if out_dir.exists():
        raise FileExistsError(f"refusing to overwrite protected Stage-1 checkpoint: {out_dir}")
    if tmp_dir.exists():
        import shutil
        shutil.rmtree(tmp_dir, ignore_errors=True)
    tmp_dir.mkdir(parents=True, exist_ok=True)
    ring = getattr(model, "_rss_ring", None)
    if kind == "adapter":
        model.save_pretrained(str(tmp_dir / "adapter"), safe_serialization=True)
    elif ring is not None:
        ring.export(tmp_dir)
    else:
        model.to(torch.bfloat16)
        model.save_pretrained(str(tmp_dir), safe_serialization=True)
    tokenizer.save_pretrained(str(tmp_dir))
    (tmp_dir / "rss_meta.json").write_text(json.dumps({
        "stage": 1, "kind": kind, "base_model": cfg["base_model_repo"],
        "adapter_dir": "adapter" if kind == "adapter" else None, "created": time.time()}, indent=1), encoding="utf-8")
    out_dir.parent.mkdir(parents=True, exist_ok=True)
    import os
    os.replace(tmp_dir, out_dir)
    print(f"   [stage1] checkpoint atomically saved -> {out_dir} ({kind})", flush=True)
    return out_dir


# --------------------------------------------------------------- the stage
def run_pretrain(model, tokenizer, plan: Dict, cfg: Dict, surgery: Dict, pretrain_path: Path,
                 checkpoint_dir: Path) -> Dict:
    tune = cfg["tune"]
    gpu = plan["tier"] == 1
    seq_len = tune.get("pretrain_seq_len") or (512 if gpu else 192)
    max_chunks = tune.get("pretrain_max_chunks") or (2000 if gpu else 96)
    batch_size = tune.get("pretrain_batch_size") or (4 if gpu else 2)
    lr = tune.get("pretrain_lr") or (2e-4 if gpu else 2e-5)
    chunks = build_chunks(tokenizer, pretrain_path, seq_len, max_chunks, cfg.get("conv_ratio", 0.25))
    if not chunks:
        raise RuntimeError("no pre-training chunks could be built from pretrain_clean.jsonl")
    amp, scaler = make_amp(plan)
    model.config.use_cache = False
    device = model_device(model)
    ring = getattr(model, "_rss_ring", None)
    batches = batch_stream(chunks, batch_size, device)

    if gpu:
        model = attach_lora(model, plan, int(tune.get("lora_rank", 16)), True, surgery.get("resized", False))
        steps = int(tune.get("pretrain_steps") or min(300, max(20, len(chunks) // batch_size)))
        params = [p for p in model.parameters() if p.requires_grad]
        print(f"[Pretrain] Stage 1 (rsLoRA, 4-bit) for {steps} steps")
        loss = train_loop(model, params, batches, lm_loss, steps, lr, 1, amp, scaler, "cpt-lora", log_every=5)
        save_stage1_checkpoint(model, tokenizer, checkpoint_dir, cfg, "adapter")
        return {"kind": "adapter", "loss": loss, "checkpoint": str(checkpoint_dir)}

    # ---------------- CPU elevator ----------------
    layers = find_layers(model)
    total_layers = len(layers)
    window = 2
    steps_per_window = int(tune.get("pretrain_steps_per_window") or 4)
    base_dtype = next(model.parameters()).dtype
    for p in model.parameters():
        p.requires_grad_(False)
    if ring is not None:
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    print(f"[Pretrain] Stage 1 elevator: {total_layers} layers, window={window}, {steps_per_window} steps/window, "
          f"seq={seq_len} batch={batch_size} lr={lr}")
    window_losses = []
    for start in range(0, total_layers, window):
        idxs = list(range(start, min(total_layers, start + window)))
        if ring is not None:
            ring.pin(idxs)
        params: List[torch.nn.Parameter] = []
        for i in idxs:
            if base_dtype != torch.float32:
                layers[i].to(torch.float32)
            for p in layers[i].parameters():
                p.requires_grad_(True)
                params.append(p)
        print(f"   [elevator] floor {idxs[0]}-{idxs[-1]} unfrozen "
              f"({sum(p.numel() for p in params) / 1e6:.1f}M params)" + (f" | {ring.describe()}" if ring else ""))
        loss = train_loop(model, params, batches, lm_loss, steps_per_window, lr, 1, amp, scaler,
                          f"elevator {idxs[0]}-{idxs[-1]}", log_every=max(1, steps_per_window))
        window_losses.append(loss)
        for i in idxs:
            for p in layers[i].parameters():
                p.requires_grad_(False)
            if base_dtype != torch.float32:
                layers[i].to(base_dtype)
        if ring is not None:
            ring.unpin(idxs, save=True)
        del params
        gc.collect()
    if ring is not None:
        model.gradient_checkpointing_disable()
    model.eval()
    save_stage1_checkpoint(model, tokenizer, checkpoint_dir, cfg, "full")
    return {"kind": "full", "loss": sum(window_losses) / max(len(window_losses), 1),
            "checkpoint": str(checkpoint_dir)}