from __future__ import annotations
import argparse
import json
import os
import time
import torch
import torch.nn as nn
from safetensors.torch import save_file
from data import assemble, licenses
from data.bio import PAD_ID, encode_with_offsets
from data.taxonomy import JIGSAW_LABELS, N_PII_TAGS, TAXONOMY
from modeling import FELAModerationV2, ModerationConfig
from fela_server.train import TrainCtx, chunk_train


def encode_depth(model, ids, mask, depth):
    x = model.tok_emb(ids)
    kpm = mask == 0
    for blk in model.blocks[:depth]:
        x = blk(x, kpm)
    return model.norm(x)


def head_logits(model, x, mask, task):
    if task == "pii":
        return model.pii_head(x)
    pooled = model._pool(x, mask)
    heads = {"taxonomy": model.tax_head, "jigsaw": model.jig_head, **model._clf_tasks}
    return heads[task](pooled)


def masked_bce(bce, logits, targets, tmask, posw=1.0):
    loss = bce(logits, targets)
    w = tmask * (1.0 + (posw - 1.0) * targets)
    return (loss * w).sum() / w.sum().clamp(min=1.0)


def _ckpt_uri_parts(uri):
    b, _, k = uri[5:].partition("/")
    return (b, k)


def ckpt_save(model, uri):
    import io
    import boto3

    buf = io.BytesIO()
    torch.save({k: v.cpu() for k, v in model.state_dict().items()}, buf)
    b, k = _ckpt_uri_parts(uri)
    boto3.client("s3").put_object(Bucket=b, Key=k, Body=buf.getvalue())


def ckpt_load(model, uri):
    import io
    import boto3
    from botocore.exceptions import ClientError

    b, k = _ckpt_uri_parts(uri)
    try:
        body = boto3.client("s3").get_object(Bucket=b, Key=k)["Body"].read()
    except ClientError:
        return False
    model.load_state_dict(torch.load(io.BytesIO(body), map_location="cpu"))
    return True


def make_tox_batch(examples, max_len):
    seqs = [encode_with_offsets(ex["text"], max_len)[0] for ex in examples]
    length = min(max((len(s) for s in seqs)), max_len)
    b = len(seqs)
    ids = torch.full((b, length), PAD_ID, dtype=torch.long)
    am = torch.zeros((b, length), dtype=torch.long)
    for i, s in enumerate(seqs):
        n = min(len(s), length)
        ids[i, :n] = torch.tensor(s[:n])
        am[i, :n] = 1
    labels = torch.tensor([ex["labels"] for ex in examples], dtype=torch.float)
    lmask = torch.tensor([ex["mask"] for ex in examples], dtype=torch.float)
    return (ids, am, labels, lmask)


def iter_tox_batches(loader_fn, batch_size, max_len, cap):
    yield from iter_clf_batches(loader_fn(max_rows=cap), batch_size, max_len)


def iter_clf_batches(stream, batch_size, max_len):
    buf = []
    for ex in stream:
        buf.append(ex)
        if len(buf) >= batch_size:
            yield make_tox_batch(buf, max_len)
            buf = []
    if buf:
        yield make_tox_batch(buf, max_len)


def _combined(real_fn, synth_fn, real_cap, n_synth, seed=0):
    import random

    rng = random.Random(seed)
    pools = []
    if real_fn is not None:
        pools.append(real_fn(max_rows=real_cap))
    if synth_fn is not None:
        pools.append(synth_fn(n_synth, seed=seed))
    while pools:
        src = rng.choice(pools)
        try:
            yield next(src)
        except StopIteration:
            pools.remove(src)


_AUX_HEADS = [
    ("spam", "spam", "load_spam_examples", "synth_spam"),
    ("jailbreak", "jailbreak", "load_jailbreak_examples", "synth_jailbreak"),
    ("nsfw", "nsfw", "load_nsfw_examples", "synth_nsfw"),
    ("identity", "identity", "load_identity_examples", None),
]


def quantize_int8(state):
    out, scales = ({}, {})
    for k, v in state.items():
        if v.dtype.is_floating_point and v.dim() >= 2:
            amax = v.abs().max().clamp(min=1e-08)
            s = amax / 127.0
            out[k] = (v / s).round().clamp(-127, 127).to(torch.int8)
            scales[k] = float(s)
        else:
            out[k] = v
    return (out, scales)


def subset_state(state, depth):
    keep = {}
    for k, v in state.items():
        if k.startswith("blocks."):
            if int(k.split(".")[1]) < depth:
                keep[k] = v
        else:
            keep[k] = v
    return keep


def _bytes_int8(state):
    return sum(
        (
            v.numel() * (1 if v.dtype == torch.int8 else v.element_size())
            for v in state.values()
        )
    )


def export(model, out_dir, n_tax, nested_depth):
    os.makedirs(out_dir, exist_ok=True)
    state = {k: v.detach().cpu().float() for k, v in model.state_dict().items()}
    save_file(state, os.path.join(out_dir, "model.safetensors"))
    manifest = {
        "n_tax": n_tax,
        "n_pii_tags": N_PII_TAGS,
        "taxonomy": TAXONOMY,
        "jigsaw": JIGSAW_LABELS,
        "ship": "weights_only",
        "tiers": [],
    }

    def write_tier(name, st):
        q, sc = quantize_int8(st)
        save_file(
            {k: v for k, v in q.items()},
            os.path.join(out_dir, f"tier_{name}_int8.safetensors"),
        )
        with open(os.path.join(out_dir, f"tier_{name}_scales.json"), "w") as f:
            json.dump(sc, f)
        manifest["tiers"].append(
            {
                "tier": name,
                "int8_bytes": _bytes_int8(q),
                "int8_mb": round(_bytes_int8(q) / 1000000.0, 2),
            }
        )

    if 0 < nested_depth < len(model.blocks):
        write_tier("small", subset_state(state, nested_depth))
    write_tier("full", state)
    with open(os.path.join(out_dir, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)
    with open(os.path.join(out_dir, "NOTICE.txt"), "w") as f:
        f.write(licenses.notice_text())
    return manifest


def _saver_for(n_tax, nested_depth):

    def _save(model, artifact, **meta):
        manifest = export(model, artifact or "./retrained", n_tax, nested_depth)
        print("EXPORT:", json.dumps(manifest["tiers"]))
        if meta:
            print("METRICS:", json.dumps({k: v for k, v in meta.items()}, default=str))

    return _save


def _train_core(ctx, args):
    torch.manual_seed(ctx.rank)
    torch.set_num_threads(
        int(os.environ.get("OMP_NUM_THREADS", torch.get_num_threads()))
    )
    cfg = ModerationConfig(n_pii_tags=N_PII_TAGS)
    model = FELAModerationV2(cfg, n_tax=len(TAXONOMY))
    if ctx.rank == 0:
        print(
            f"params={model.param_count() / 1000000.0:.2f}M n_tax={len(TAXONOMY)} pii_tags={N_PII_TAGS} world={ctx.world} threads={torch.get_num_threads()}"
        )
    resume_uri = os.environ.get("MOD_RESUME_S3")
    if resume_uri and ckpt_load(model, resume_uri):
        print(f"[rank{ctx.rank}] resumed from {resume_uri}")
    ckpt_uri = os.environ.get("MOD_CKPT_S3")
    ckpt_every = int(os.environ.get("MOD_CKPT_EVERY", "300"))
    posw = float(os.environ.get("MOD_POSWEIGHT", "5.0"))
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
    bce = nn.BCEWithLogitsLoss(reduction="none")
    ce = nn.CrossEntropyLoss(ignore_index=-100)
    depths = [len(model.blocks)] + (
        [args.nested_depth] if 0 < args.nested_depth < len(model.blocks) else []
    )
    use_tox = not args.smoke and (not args.no_toxicity)
    core = []
    aux = []
    if use_tox:
        import importlib
        from data import toxicity

        ml = max(args.context_lengths)
        core.append(
            (
                "taxonomy",
                iter_tox_batches(
                    toxicity.load_taxonomy_examples, args.bs, ml, args.tox_cap
                ),
            )
        )
        core.append(
            (
                "jigsaw",
                iter_tox_batches(
                    toxicity.load_jigsaw_examples, args.bs, ml, args.tox_cap
                ),
            )
        )
        n_syn = max(1, args.n_synth // ctx.world)
        for task, modname, real_attr, synth_attr in _AUX_HEADS:
            try:
                mod = importlib.import_module(f"data.{modname}")
                real_fn = getattr(mod, real_attr, None)
                synth_fn = getattr(mod, synth_attr, None) if synth_attr else None
                stream = _combined(
                    real_fn, synth_fn, args.tox_cap, n_syn, seed=ctx.rank
                )
                aux.append((task, iter_clf_batches(stream, args.bs, ml)))
            except Exception as e:
                print(f"[aux] {task} unavailable: {e!r}")
    step, toks, last, t0 = (0, 0, float("nan"), time.time())
    model.train()
    for ep in range(args.epochs):
        pii = assemble.iter_pii_batches(
            args.bs,
            args.context_lengths,
            max(1, args.n_synth // ctx.world),
            seed=ep * ctx.world + ctx.rank,
            use_real=not args.smoke and args.use_real,
        )
        for ids, mask, tags in pii:
            loss = torch.zeros(())
            for d in depths:
                x = encode_depth(model, ids, mask, d)
                loss = loss + ce(
                    head_logits(model, x, mask, "pii").reshape(-1, N_PII_TAGS),
                    tags.reshape(-1),
                )
            for cname, cit in core:
                cbatch = next(cit, None)
                if cbatch is not None:
                    cids, cam, clab, cmsk = cbatch
                    cx = encode_depth(model, cids, cam, len(model.blocks))
                    loss = loss + masked_bce(
                        bce, head_logits(model, cx, cam, cname), clab, cmsk, posw
                    )
            if aux:
                name, it = aux[step % len(aux)]
                batch = next(it, None)
                if batch is not None:
                    tids, tam, tlab, tmsk = batch
                    tx = encode_depth(model, tids, tam, len(model.blocks))
                    loss = loss + masked_bce(
                        bce, head_logits(model, tx, tam, name), tlab, tmsk, posw
                    )
            opt.zero_grad()
            loss.backward()
            opt.step()
            ctx.step(model)
            step += 1
            toks += ids.numel()
            last = float(loss.item())
            if ctx.rank == 0 and (step % 20 == 0 or args.smoke):
                dt = time.time() - t0
                print(
                    f"ep{ep} step{step} loss={last:.4f} tok/s={toks / max(dt, 1e-09):,.0f}"
                )
            if ctx.rank == 0 and ckpt_uri and (step % ckpt_every == 0):
                ckpt_save(model, ckpt_uri)
            if args.steps and step >= args.steps:
                break
        if args.steps and step >= args.steps:
            break
    dt = time.time() - t0
    ctx.log(
        steps=step,
        final_loss=last,
        wall_s=round(dt, 1),
        tok_per_s=round(toks / max(dt, 1e-09), 1),
    )
    ctx.save(model)
    if ctx.rank == 0:
        print("licenses:", [s[0] for s in licenses.SOURCES])
        print(licenses.HUMAN_SIGNOFF)
    return {
        "steps": step,
        "final_loss": last,
        "tok_per_s": round(toks / max(dt, 1e-09), 1),
    }


def _hp_from_env():
    e = os.environ.get
    return argparse.Namespace(
        epochs=int(e("MOD_EPOCHS", "6")),
        steps=int(e("MOD_STEPS", "0")),
        bs=int(e("MOD_BS", "32")),
        lr=float(e("MOD_LR", "3e-4")),
        context_lengths=[
            int(x) for x in e("MOD_CTXLENS", "128,512,2048,8192").split(",")
        ],
        nested_depth=int(e("MOD_NESTED", "4")),
        n_synth=int(e("MOD_NSYNTH", "50000")),
        tox_cap=int(e("MOD_TOXCAP", "200000")),
        use_real=e("MOD_USEREAL", "1") == "1",
        no_toxicity=e("MOD_NOTOX", "0") == "1",
        smoke=e("MOD_SMOKE", "0") == "1",
    )


@chunk_train(
    name="moderator",
    sync_every=int(os.environ.get("MOD_SYNC", "50")),
    artifact=os.environ.get("MOD_ARTIFACT", "./retrained"),
    saver=_saver_for(len(TAXONOMY), int(os.environ.get("MOD_NESTED", "4"))),
)
def train_job(ctx: TrainCtx = TrainCtx()):
    return _train_core(ctx, _hp_from_env())


def train(args):
    ctx = TrainCtx(
        0,
        1,
        sync_every=1,
        artifact=args.out,
        saver=_saver_for(len(TAXONOMY), args.nested_depth),
    )
    _train_core(ctx, args)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="./retrained")
    ap.add_argument("--epochs", type=int, default=6)
    ap.add_argument(
        "--steps", type=int, default=0, help="hard cap on optimizer steps (0=off)"
    )
    ap.add_argument("--bs", type=int, default=32)
    ap.add_argument("--lr", type=float, default=0.0003)
    ap.add_argument(
        "--context-lengths", type=int, nargs="+", default=[128, 512, 2048, 8192]
    )
    ap.add_argument(
        "--nested-depth", type=int, default=4, help="small-tier block depth (0=off)"
    )
    ap.add_argument("--n-synth", type=int, default=50000)
    ap.add_argument("--tox-cap", type=int, default=200000)
    ap.add_argument(
        "--use-real",
        action="store_true",
        help="include real PII datasets (needs network)",
    )
    ap.add_argument("--no-toxicity", action="store_true")
    ap.add_argument("--smoke", action="store_true", help="offline end-to-end check")
    args = ap.parse_args()
    if args.smoke:
        args.epochs, args.steps, args.n_synth = (1, 6, 200)
        args.bs, args.context_lengths = (4, [128, 512])
    train(args)