"""
Adversarial Paraphrasing (detector-guided decoding) for the Parrot paraphraser.

Technique from "Adversarial Paraphrasing: A Universal Attack for Humanizing AI-Generated Text" (Cheng et al., 2025,
arXiv:2506.07001): while the paraphraser writes, the AI detector chooses the words. At each decoding step the
paraphraser proposes its most likely next tokens (top-k inside top-p); every proposal is appended to the text written
so far and scored by the detector, and the proposal that reads as the MOST HUMAN is kept. The paraphraser keeps the
sentence fluent and on-meaning (only its own likely tokens are allowed); the detector steers word choice.

Here it adds one extra candidate per sentence next to Parrot's sampled versions. That candidate goes through exactly
the same checks as every other version (terms / numbers / citations restored and guarded, meaning, fluency, clarity
gate, Groq check), so a damaged adversarial sentence is simply not chosen.

Cost: one detector batch per guided step (sentences x candidates texts). On a dedicated GPU this is roughly 2-3x the
time of normal paraphrasing for the sentences it is used on; on CPU it is much slower, so by default it runs only in
the "100% AI-written" / report-guided modes and on at most ADV_MAX sentence pieces per job.

Space variables
  HUM_ADV        ai (default: 100% AI mode and report-guided rewriting) | always | off
  ADV_K          candidate tokens scored per step (default 8; the paper uses 50)
  ADV_TOP_P      nucleus for the candidates (default 0.95)
  ADV_MIN_P      a candidate needs at least this share of the top token's probability (default 0.05): keeps the
                 detector from picking an unlikely, ungrammatical word
  ADV_EVERY      guide every n-th token (default 1 = every token, as in the paper; 2 halves the detector cost)
  ADV_BATCH      sentences decoded together (default 16)
  ADV_MAX        most sentence pieces per job (default 400)
  ADV_ON_ZERO    1 = also run on ZeroGPU hardware (on the CPU copy of the model; slow). Default 0.
"""
import os

MODE = os.getenv("HUM_ADV", "ai").strip().lower()
K = max(2, int(os.getenv("ADV_K", "8")))
TOP_P = float(os.getenv("ADV_TOP_P", "0.95"))
MIN_P = float(os.getenv("ADV_MIN_P", "0.05"))
EVERY = max(1, int(os.getenv("ADV_EVERY", "1")))
BATCH = max(1, int(os.getenv("ADV_BATCH", "16")))
MAX_PIECES = max(0, int(os.getenv("ADV_MAX", "400")))
ON_ZERO = os.getenv("ADV_ON_ZERO", "0") == "1"
STATS = {"runs": 0, "pieces": 0, "steps": 0, "scored": 0, "errors": 0, "last_error": ""}


def wanted(full_ai):
    """Is detector-guided decoding used for this job?"""
    if MODE in ("off", "0", "no", "false"):
        return False
    return MODE == "always" or bool(full_ai)


def _no_repeat(prev_ids, n=3):
    """Tokens that would repeat an n-gram already written (same rule as Parrot's no_repeat_ngram_size=3)."""
    if len(prev_ids) < n - 1:
        return set()
    tail = tuple(prev_ids[-(n - 1):])
    banned = set()
    for i in range(len(prev_ids) - n + 1):
        if tuple(prev_ids[i:i + n - 1]) == tail:
            banned.add(prev_ids[i + n - 1])
    return banned


def generate(texts, model, tokenizer, device, score_fn, temperature=1.0, prefix="paraphrase: "):
    """Detector-guided paraphrase of each text (one output per text; the input text where decoding fails).
    score_fn(list[str]) -> list[float] AI probabilities (lower = more human), or None."""
    import torch
    outs = list(texts)
    if not texts:
        return outs
    STATS["runs"] += 1
    temperature = max(0.3, float(temperature))
    order = sorted(range(len(texts)), key=lambda i: len(texts[i]))      # similar lengths decode together
    for b0 in range(0, len(order), BATCH):
        ids_ = order[b0:b0 + BATCH]
        batch = [texts[i] for i in ids_]
        try:
            res = _generate_batch(batch, model, tokenizer, device, score_fn, temperature, prefix, torch)
        except Exception as err:                    # never break the job: Parrot's normal versions remain
            STATS["errors"] += 1
            STATS["last_error"] = str(err)[:200]
            print(f"[adversarial] batch skipped: {err}", flush=True)
            continue
        for i, r in zip(ids_, res):
            if r and r.strip():
                outs[i] = r.strip()
        STATS["pieces"] += len(batch)
    return outs


def _generate_batch(batch, model, tokenizer, device, score_fn, temperature, prefix, torch):
    enc = tokenizer([prefix + t for t in batch], return_tensors="pt", padding=True, truncation=True, max_length=256)
    enc = {k: v.to(device) for k, v in enc.items()}
    n = len(batch)
    max_new = min(180, int(enc["input_ids"].shape[1] * 1.35) + 12)
    eos = tokenizer.eos_token_id
    start = model.config.decoder_start_token_id
    if start is None:
        start = tokenizer.pad_token_id
    written = [[] for _ in range(n)]
    done = [False] * n
    with torch.inference_mode():
        enc_out = model.get_encoder()(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"])
        last = torch.full((n, 1), start, dtype=torch.long, device=device)
        past = None
        for step in range(max_new):
            out = model(encoder_outputs=enc_out, attention_mask=enc["attention_mask"], decoder_input_ids=last,
                        past_key_values=past, use_cache=True)
            past = out.past_key_values
            probs = torch.softmax(out.logits[:, -1, :].float() / temperature, dim=-1)
            top_p, top_i = probs.topk(K, dim=-1)
            top_p, top_i = top_p.cpu().tolist(), top_i.cpu().tolist()
            choice = [eos] * n
            texts, owners = [], []
            for b in range(n):
                if done[b]:
                    continue
                banned = _no_repeat(written[b])
                cands, cum = [], 0.0
                for p, t in zip(top_p[b], top_i[b]):
                    if t in banned:
                        continue
                    if cands and (cum >= TOP_P or p < MIN_P * top_p[b][0]):
                        break
                    cands.append((t, p))
                    cum += p
                if not cands:
                    cands = [(top_i[b][0], top_p[b][0])]
                choice[b] = cands[0][0]                    # most likely token by default
                if len(cands) > 1 and step % EVERY == 0:
                    for t, _ in cands:
                        if t == eos:                           # the sentence so far, as it would end here
                            texts.append(tokenizer.decode(written[b], skip_special_tokens=True))
                        else:
                            texts.append(tokenizer.decode(written[b] + [t], skip_special_tokens=True))
                        owners.append((b, t))
            if texts:
                sc = score_fn([t if t.strip() else "." for t in texts])
                STATS["scored"] += len(texts)
                if sc:
                    best = {}
                    for (b, t), s in zip(owners, sc):
                        if b not in best or s < best[b][0] - 1e-6:   # ties keep the more likely (earlier) token
                            best[b] = (s, t)
                    for b, (_, t) in best.items():
                        choice[b] = t
            STATS["steps"] += 1
            for b in range(n):
                if done[b]:
                    choice[b] = tokenizer.pad_token_id
                    continue
                if choice[b] == eos:
                    done[b] = True
                else:
                    written[b].append(choice[b])
            if all(done):
                break
            last = torch.tensor([[c] for c in choice], dtype=torch.long, device=device)
    return [tokenizer.decode(w, skip_special_tokens=True) for w in written]


def status():
    if not wanted(True):
        return "Adversarial paraphrasing: off (HUM_ADV=off)."
    s = STATS
    return (f"Adversarial paraphrasing ({MODE}, k={K}, every {EVERY} token): {s['pieces']} sentence pieces, "
            f"{s['scored']:,} detector checks" + (f"; last error: {s['last_error']}" if s["last_error"] else "."))