"""Stage 2 rational SFT with checkpoint-safe active-vocabulary loss slicing.

The Stage-1 embedding and LM-head tensors are never resized or rewritten here. Inputs retain
original tokenizer IDs; only the output projection and CE classes are sliced to ACTIVE_TOKEN_IDS.
"""
from __future__ import annotations

import gc
import itertools
import json
import math
import os
import shutil
import time
from datetime import datetime
from pathlib import Path
from typing import Dict, Iterable, List, Optional

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

from .elevator_pretrainer import attach_lora, make_amp, model_device
from .tokenizer_surgeon import (build_active_vocab_ids, render_prompt, render_target,
                                summarize_active_vocab)


def _memory_snapshot() -> str:
    rss = None
    try:
        import psutil
        rss = psutil.Process(os.getpid()).memory_info().rss
    except Exception:
        pass
    cgroup = None
    for candidate in ("/sys/fs/cgroup/memory.current", "/sys/fs/cgroup/memory/usage_in_bytes"):
        try:
            cgroup = int(Path(candidate).read_text().strip())
            break
        except Exception:
            continue
    fields = []
    if rss is not None:
        fields.append(f"rss={rss / (1024 ** 2):.1f}MiB")
    if cgroup is not None:
        fields.append(f"cgroup={cgroup / (1024 ** 2):.1f}MiB")
    return " ".join(fields) if fields else "memory=n/a"


def _diag(label: str, started: Optional[float] = None, detail: str = "") -> None:
    elapsed = f" dt={time.perf_counter() - started:.3f}s" if started is not None else ""
    extra = f" {detail}" if detail else ""
    print(f"[SFT-DIAG] {datetime.now().astimezone().isoformat(timespec='seconds')} "
          f"{label}{elapsed} {_memory_snapshot()}{extra}", flush=True)


# -------------------------------------------------------------------- data
def build_examples(tokenizer, path: Path, max_len: int, im_end_id: int,
                   max_thinking_tokens: Optional[int] = None) -> List[Dict]:
    """Tokenise every valid SFT record; trim only prompt context, never discard an example."""
    examples: List[Dict] = []
    modes: Dict[str, int] = {}
    prompt_trimmed = 0
    long_targets = 0
    dataset_token_ids = set()
    with open(path, "r", encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, 1):
            if not line.strip():
                continue
            rec = json.loads(line)
            thinking = rec["thinking"]
            original_prompt = render_prompt(rec["system"], rec["user"], rec["mode"])
            original_target = render_target(thinking, rec["assistant"], rec["mode"])
            full_p_ids = tokenizer(original_prompt, add_special_tokens=False)["input_ids"]
            full_t_ids = tokenizer(original_target, add_special_tokens=False)["input_ids"]
            # Preserve all original dataset IDs before any prompt/thinking budget is applied.
            dataset_token_ids.update(int(x) for x in full_p_ids)
            dataset_token_ids.update(int(x) for x in full_t_ids)
            if max_thinking_tokens is not None and rec["mode"] != "non_thinking":
                thinking_ids = tokenizer(thinking, add_special_tokens=False)["input_ids"]
                if len(thinking_ids) > max_thinking_tokens:
                    thinking = tokenizer.decode(thinking_ids[:max_thinking_tokens], skip_special_tokens=False)
                    capped_target = render_target(thinking, rec["assistant"], rec["mode"])
                    full_t_ids = tokenizer(capped_target, add_special_tokens=False)["input_ids"]
                    dataset_token_ids.update(int(x) for x in full_t_ids)
            if not full_t_ids or int(full_t_ids[-1]) != int(im_end_id):
                full_t_ids.append(int(im_end_id))
                dataset_token_ids.add(int(im_end_id))

            p_ids, t_ids = list(full_p_ids), list(full_t_ids)
            if len(p_ids) + len(t_ids) > max_len:
                prompt_budget = max(0, max_len - len(t_ids))
                p_ids = p_ids[-prompt_budget:] if prompt_budget else []
                prompt_trimmed += 1
            if len(t_ids) > max_len:
                # The answer/thinking target is never dropped or truncated; correctness takes precedence.
                long_targets += 1
            examples.append({"input_ids": p_ids + t_ids,
                             "labels": [-100] * len(p_ids) + t_ids,
                             "source_line": line_no})
            modes[rec["mode"]] = modes.get(rec["mode"], 0) + 1

    # Carry original dataset IDs for the vocab allow-list, even if prompt clipping occurred.
    if examples:
        examples[0]["_dataset_token_ids"] = sorted(dataset_token_ids)
    print(f"   [sft-data] {len(examples)} examples retained ({', '.join(f'{k}={v}' for k, v in sorted(modes.items()))}); "
          f"prompt-context-trimmed={prompt_trimmed}, long-targets-kept={long_targets}", flush=True)
    return examples


def _collate_sft(rows: List[Dict], pad_id: int) -> Dict[str, torch.Tensor]:
    longest = max(len(row["input_ids"]) for row in rows)
    ids = torch.full((len(rows), longest), int(pad_id), dtype=torch.long)
    labels = torch.full((len(rows), longest), -100, dtype=torch.long)
    mask = torch.zeros((len(rows), longest), dtype=torch.long)
    for i, row in enumerate(rows):
        n = len(row["input_ids"])
        ids[i, :n] = torch.as_tensor(row["input_ids"], dtype=torch.long)
        labels[i, :n] = torch.as_tensor(row["labels"], dtype=torch.long)
        mask[i, :n] = 1
    return {"input_ids": ids, "attention_mask": mask, "labels": labels}


def _epoch_batch_stream(examples: List[Dict], epochs: int, pad_id: int, device):
    # Construct the DataLoader immediately (not on first generator iteration) so diagnostics
    # bracket the real initialization. num_workers=0 avoids multiprocessing semaphores.
    generator = torch.Generator()
    generator.manual_seed(11)
    loader = DataLoader(examples, batch_size=1, shuffle=True, generator=generator,
                        num_workers=0, pin_memory=False,
                        collate_fn=lambda rows: _collate_sft(rows, pad_id))
    def stream():
        for _epoch in range(epochs):
            for batch in loader:
                yield {key: value.to(device, non_blocking=False) for key, value in batch.items()}
    return stream()


# ---------------------------------------------------------------- sliced output/loss
def _base_model_and_head(model):
    base = model.get_base_model() if hasattr(model, "get_base_model") else model
    transformer = getattr(base, "model", None)
    if transformer is None:
        raise RuntimeError("active-vocab slicing needs a causal model exposing .model (transformer backbone)")
    head = base.get_output_embeddings() if hasattr(base, "get_output_embeddings") else getattr(base, "lm_head", None)
    # PEFT can wrap modules_to_save in a ModulesToSaveWrapper. Resolve the currently active
    # saved module (or its original module) rather than assuming wrapper.weight exists.
    if head is not None and not hasattr(head, "weight"):
        saved = getattr(head, "modules_to_save", None)
        active = getattr(head, "active_adapter", None)
        active = active[0] if isinstance(active, (list, tuple)) and active else active
        if saved is not None and active in saved:
            head = saved[active]
        elif getattr(head, "original_module", None) is not None:
            head = head.original_module
    if head is None or not hasattr(head, "weight"):
        raise RuntimeError("active-vocab slicing could not locate the output lm_head weight")
    return base, transformer, head


def make_stop_weighted_loss(im_end_id: int, stop_weight: float, active_token_ids: List[int],
                            diagnostic=_diag, ring=None):
    """Use only active lm_head rows and map original IDs to compact CE classes."""
    active_ids_cpu = torch.as_tensor(active_token_ids, dtype=torch.long, device="cpu")
    cached_head = None
    cached_bias = None
    vocab_map = None
    compact_stop_id = None
    output_rows = None

    def loss_fn(model, batch):
        nonlocal cached_head, cached_bias, vocab_map, compact_stop_id, output_rows
        base, transformer, head = _base_model_and_head(model)
        ring_handle = ring or getattr(base, "_rss_ring", None) or getattr(model, "_rss_ring", None)
        if cached_head is None:
            started = time.perf_counter()
            diagnostic("HEAD-SLICE-BEGIN", detail=f"active_ids={active_ids_cpu.numel()}")
            if ring_handle is not None and "head" not in getattr(ring_handle, "loaded", set()):
                ring_handle._ensure("head")
            weight = head.weight
            output_rows = int(weight.shape[0])
            if int(active_ids_cpu[-1]) >= output_rows:
                raise RuntimeError(f"active token ID {int(active_ids_cpu[-1])} exceeds lm_head rows {output_rows}")
            index = active_ids_cpu.to(weight.device)
            with torch.no_grad():
                cached_head = weight.index_select(0, index).detach()
                bias = getattr(head, "bias", None)
                cached_bias = bias.index_select(0, index).detach() if bias is not None else None
                vocab_map = torch.full((output_rows,), -1, dtype=torch.long, device=weight.device)
                vocab_map[index] = torch.arange(index.numel(), dtype=torch.long, device=weight.device)
                compact_stop_id = int(vocab_map[int(im_end_id)].item()) if int(im_end_id) < output_rows else -1
            if ring_handle is not None and getattr(ring_handle, "evict_ends", False):
                ring_handle._evict("head")
            if compact_stop_id < 0:
                raise RuntimeError(f"stop token {im_end_id} is not in ACTIVE_TOKEN_IDS")
            diagnostic("HEAD-SLICE-DONE", started, detail=f"head_rows={output_rows} active={len(index)}")

        t0 = time.perf_counter()
        diagnostic("FORWARD-BEGIN")
        outputs = transformer(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"],
                              use_cache=False, return_dict=True)
        hidden = outputs.last_hidden_state if hasattr(outputs, "last_hidden_state") else outputs[0]
        diagnostic("FORWARD-DONE", t0, detail=f"hidden={tuple(hidden.shape)}")

        t0 = time.perf_counter()
        diagnostic("LOSS-BEGIN")
        sliced_logits = F.linear(hidden[:, :-1, :], cached_head, cached_bias)
        labels = batch["labels"][:, 1:].to(device=sliced_logits.device, dtype=torch.long)
        valid = labels.ne(-100)
        safe_labels = labels.clamp(min=0, max=output_rows - 1)
        compact_labels = vocab_map[safe_labels]
        missing = valid & compact_labels.lt(0)
        if bool(missing.any()):
            bad_ids = torch.unique(labels[missing]).detach().cpu().tolist()[:12]
            raise RuntimeError(f"SFT labels missing from ACTIVE_TOKEN_IDS (first IDs: {bad_ids})")
        compact_labels = compact_labels.masked_fill(~valid, -100)
        logits_for_ce = sliced_logits.float() if sliced_logits.dtype in (torch.float16, torch.bfloat16) else sliced_logits
        per_token = F.cross_entropy(logits_for_ce.reshape(-1, logits_for_ce.shape[-1]),
                                    compact_labels.reshape(-1), ignore_index=-100,
                                    reduction="none").view_as(compact_labels)
        weights = valid.to(dtype=per_token.dtype)
        is_stop = compact_labels.eq(compact_stop_id) & valid
        weights.mul_(1.0 + (float(stop_weight) - 1.0) * is_stop.to(dtype=weights.dtype))
        loss = (per_token * weights).sum() / weights.sum().clamp_min(1.0)
        diagnostic("LOSS-DONE", t0, detail=f"sliced_logits={tuple(sliced_logits.shape)} loss={float(loss.detach()):.5f}")
        return loss
    return loss_fn


# ------------------------------------------------------------------ optimiser loop
def _train_sft_loop(model, params: List[torch.nn.Parameter], batches: Iterable[Dict], steps: int,
                    lr: float, accum: int, amp, tag: str = "sft") -> float:
    _diag("OPTIMIZER-INIT-BEGIN", detail=f"trainable_params={sum(p.numel() for p in params):,}")
    started = time.perf_counter()
    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))

    scheduler = torch.optim.lr_scheduler.LambdaLR(opt, lr_at)
    _diag("OPTIMIZER-INIT-DONE", started)
    model.train()
    batch_iter = iter(batches)
    history: List[float] = []
    for step in range(steps):
        group = list(itertools.islice(batch_iter, accum))
        if not group:
            break
        opt.zero_grad(set_to_none=True)
        step_loss = 0.0
        for micro_idx, batch in enumerate(group, 1):
            _diag("MICROBATCH-BEGIN", detail=f"step={step + 1}/{steps} micro={micro_idx}/{len(group)}")
            with amp():
                loss = batch["_loss_fn"](model, batch) / len(group)
            t0 = time.perf_counter()
            _diag("BACKWARD-BEGIN", detail=f"step={step + 1} micro={micro_idx}")
            loss.backward()
            _diag("BACKWARD-DONE", t0, detail=f"step={step + 1} micro={micro_idx}")
            step_loss += float(loss.detach())
        t0 = time.perf_counter()
        _diag("GRAD-CLIP-BEGIN", detail=f"step={step + 1}")
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        _diag("GRAD-CLIP-DONE", t0)
        t0 = time.perf_counter()
        _diag("OPTIMIZER-STEP-BEGIN", detail=f"step={step + 1}")
        opt.step()
        scheduler.step()
        opt.zero_grad(set_to_none=True)
        _diag("OPTIMIZER-STEP-DONE", t0, detail=f"lr={opt.param_groups[0]['lr']:.2e}")
        history.append(step_loss)
        print(f"   [{tag}] step {step + 1:>4}/{steps} loss {step_loss:.5f} "
              f"lr {opt.param_groups[0]['lr']:.2e}", flush=True)
    del opt, scheduler
    gc.collect()
    tail = history[len(history) // 2:] or history
    return sum(tail) / max(len(tail), 1)


# ------------------------------------------------------------------- stage
def _atomic_save_stage2(model, tokenizer, final_dir: Path, meta: Dict, active_ids: List[int], ring=None) -> None:
    final_dir = Path(final_dir)
    tmp_dir = final_dir.with_name(final_dir.name + ".tmp")
    backup_dir = final_dir.with_name(final_dir.name + ".previous")
    _diag("CHECKPOINT-SAVE-BEGIN", detail=f"tmp={tmp_dir}")
    started = time.perf_counter()
    if tmp_dir.exists():
        shutil.rmtree(tmp_dir, ignore_errors=True)
    tmp_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(tmp_dir / "adapter") if meta.get("has_adapter") else str(tmp_dir),
                          safe_serialization=True)
    if ring is not None and meta.get("has_adapter"):
        ring.export(tmp_dir)
    tokenizer.save_pretrained(str(tmp_dir))
    with open(tmp_dir / "active_vocab_ids.json", "w", encoding="utf-8") as handle:
        json.dump(active_ids, handle, separators=(",", ":"))
        handle.flush()
        os.fsync(handle.fileno())
    with open(tmp_dir / "rss_meta.json", "w", encoding="utf-8") as handle:
        json.dump(meta, handle, ensure_ascii=False, indent=1)
        handle.flush()
        os.fsync(handle.fileno())
    if backup_dir.exists():
        shutil.rmtree(backup_dir, ignore_errors=True)
    moved_old = False
    try:
        if final_dir.exists():
            os.replace(final_dir, backup_dir)
            moved_old = True
        os.replace(tmp_dir, final_dir)
    except Exception:
        if moved_old and not final_dir.exists() and backup_dir.exists():
            os.replace(backup_dir, final_dir)
        raise
    else:
        if backup_dir.exists():
            shutil.rmtree(backup_dir, ignore_errors=True)
    _diag("CHECKPOINT-SAVE-DONE", started, detail=f"dir={final_dir}")


def run_finetune(model, tokenizer, plan: Dict, cfg: Dict, surgery: Dict, sft_path: Path,
                 final_dir: Path, adapter_dir: Optional[str] = None) -> Dict:
    tune = cfg["tune"]
    gpu = plan["tier"] == 1
    ring = getattr(model, "_rss_ring", None)
    max_len = int(tune.get("sft_max_seq_len") or (1024 if gpu else 384))
    epochs = int(tune.get("sft_epochs") or 3)
    lr = float(tune.get("sft_lr") or (2e-4 if gpu else 1e-4))
    stop_weight = float(tune.get("stop_token_weight") or 3.0)
    rank = int(tune.get("lora_rank", 16))
    micro_batch = 1  # CPU-safe micro-batches; accumulation preserves the requested effective batch.
    grad_accum = int(tune.get("gradient_accumulation_steps") or (2 if not gpu else 1))
    max_thinking_tokens = tune.get("sft_max_thinking_tokens", 5)
    if max_thinking_tokens is not None:
        max_thinking_tokens = max(0, int(max_thinking_tokens))

    _diag("LORA-INIT-BEGIN", detail=f"rank={rank} adapter_dir={adapter_dir or 'new'}")
    started = time.perf_counter()
    amp, _scaler = make_amp(plan)
    model.config.use_cache = False
    if ring is not None:
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    # Do not wrap/copy full embed_tokens or lm_head in modules_to_save for SFT.
    model = attach_lora(model, plan, rank, False, surgery.get("resized", False),
                        adapter_dir=adapter_dir, save_resized_modules=False)
    for name, parameter in model.named_parameters():
        is_lora = "lora_A" in name or "lora_B" in name
        parameter.requires_grad_(is_lora)
    params = [p for p in model.parameters() if p.requires_grad]
    if not params:
        raise RuntimeError("LoRA initialised but no lora_A/lora_B parameters are trainable")
    non_lora_trainable = [name for name, p in model.named_parameters() if p.requires_grad and
                          "lora_A" not in name and "lora_B" not in name]
    if non_lora_trainable:
        raise RuntimeError(f"non-LoRA trainable parameters detected: {non_lora_trainable[:8]}")
    trainable_n = sum(p.numel() for p in params)
    device = model_device(model)
    _diag("LORA-INIT-DONE", started, detail=f"trainable={trainable_n:,} device={device}")

    _diag("SFT-DATA-INIT-BEGIN", detail=f"path={sft_path}")
    started = time.perf_counter()
    examples = build_examples(tokenizer, sft_path, max_len, surgery["ids"]["im_end"],
                              max_thinking_tokens=max_thinking_tokens)
    if not examples:
        raise RuntimeError("no SFT examples available")
    dataset_ids = examples[0].pop("_dataset_token_ids", [])
    _base, _transformer, head = _base_model_and_head(model)
    configured_rows = int(getattr(getattr(model, "config", None), "vocab_size", 0) or 0)
    full_vocab_rows = int(head.weight.shape[0]) if head.weight.numel() else max(
        configured_rows, int(getattr(tokenizer, "vocab_size", len(tokenizer))), len(tokenizer))
    active_ids = build_active_vocab_ids(tokenizer, dataset_ids, vocab_rows=full_vocab_rows)
    vocab_stats = summarize_active_vocab(tokenizer, dataset_ids, active_ids, full_vocab_rows)
    steps = max(1, math.ceil((len(examples) * epochs) / max(1, grad_accum)))
    if tune.get("sft_max_steps"):
        steps = min(steps, int(tune["sft_max_steps"]))
    _diag("ACTIVE-VOCAB-SLICE-DONE", started,
          detail=f"full={full_vocab_rows} active={len(active_ids)} inactive={full_vocab_rows-len(active_ids)} "
                 f"dataset_ids={len(dataset_ids)}")
    print(f"   [vocab-slice] active={len(active_ids):,}/{full_vocab_rows:,}; "
          f"inactive={full_vocab_rows-len(active_ids):,}; "
          f"ASCII/code={vocab_stats['english_ascii_code_math_kept']:,}; "
          f"Devanagari={vocab_stats['devanagari_kept']:,}; "
          f"special={vocab_stats['special_control_kept']:,}; bytes={vocab_stats['byte_fallback_kept']:,}; "
          f"dataset-only foreign kept={vocab_stats['dataset_only_foreign_kept']:,}; "
          f"foreign pruned={vocab_stats['foreign_vocab_pruned']:,}; "
          "input embedding/checkpoint rows unchanged", flush=True)

    _diag("DATALOADER-INIT-BEGIN", detail=f"batch_size={micro_batch} workers=0 pin_memory=False")
    started = time.perf_counter()
    batches = _epoch_batch_stream(examples, epochs, surgery["ids"]["pad"], device)
    # Force DataLoader construction now so initialization failures have an exact diagnostic marker.
    _diag("DATALOADER-INIT-DONE", started, detail="num_workers=0 pin_memory=False")
    loss_core = make_stop_weighted_loss(surgery["ids"]["im_end"], stop_weight, active_ids, ring=ring)
    def loss_with_batch(model_obj, batch):
        return loss_core(model_obj, batch)
    # Keep loss callback out of collated tensors; bind it explicitly for this local loop.
    def prepared_batches():
        for batch in batches:
            batch["_loss_fn"] = loss_with_batch
            yield batch

    print(f"[Finetune] Stage 2 rational SFT: {steps} steps (epochs={epochs}, micro-batch={micro_batch}, "
          f"grad-accum={grad_accum}, effective-batch={micro_batch * grad_accum}, lr={lr}, "
          f"stop-token weight x{stop_weight})", flush=True)
    loss = _train_sft_loop(model, params, prepared_batches(), steps, lr, grad_accum, amp)
    if ring is not None:
        model.gradient_checkpointing_disable()

    final_dir = Path(final_dir)
    meta = {"stage": 2, "base_model": cfg["base_model_repo"], "tier": plan["tier"],
            "modes": ["thinking", "non_thinking", "max_thinking"], "created": time.time(),
            "final_loss": loss, **vocab_stats,
            "active_vocab_file": "active_vocab_ids.json", "gradient_accumulation_steps": grad_accum,
            "micro_batch_size": micro_batch, "stop_token_weight": stop_weight}

    _diag("CHECKPOINT-PREP-BEGIN")
    if plan["tier"] == 2 and not plan.get("quant") and ring is None:
        model = model.merge_and_unload()
        model.to(torch.bfloat16)
        meta.update({"has_adapter": False, "adapter_dir": None, "weights": "merged-full"})
    else:
        meta.update({"has_adapter": True, "adapter_dir": "adapter",
                     "weights": "ring-base+adapter" if ring is not None else "adapter-only",
                     "base_is_self": ring is not None})
    _atomic_save_stage2(model, tokenizer, final_dir, meta, active_ids, ring=ring)
    gc.collect()
    return {"final_dir": str(final_dir), "loss": loss, "meta": meta,
            "active_vocab_size": len(active_ids), "full_vocab_size": full_vocab_rows}