"""Zap tokenization, windowing, label space and training-time augmentation."""
import random

from zap.extract import SPECIAL_TOKENS, serialize_ctx, serialize_field

FIELD_LABELS = ["username", "email", "current-password", "new-password", "one-time-code", "tel", "name",
                "given-name", "family-name", "organization", "street-address", "address-line2", "address-level2",
                "address-level1", "postal-code", "country", "cc-name", "cc-number", "cc-exp", "cc-exp-month",
                "cc-exp-year", "cc-csc", "bday", "search", "captcha", "other"]
FORM_LABELS = ["login", "signup", "password-change", "password-recovery", "otp", "payment", "address", "newsletter",
               "contact", "search", "other"]
# One token-classification label space: [CLS] carries form:*, each [FLD] carries field:*.
ID2LABEL = [f"form:{x}" for x in FORM_LABELS] + [f"field:{x}" for x in FIELD_LABELS]
LABEL2ID = {l: i for i, l in enumerate(ID2LABEL)}
N_FORM = len(FORM_LABELS)
MAX_LEN = 512


def windows(tok, ctx, fields, max_len=MAX_LEN):
    """Split a form into model inputs: each window = [CLS] ctx + as many whole fields as fit + [SEP].

    Returns a list of (input_ids, field_positions, field_indices). Shared by training and inference;
    the JS/extension port must do the same.
    """
    ctx_ids = tok(serialize_ctx(ctx), add_special_tokens=False)["input_ids"][:160]
    f_ids = [tok(serialize_field(f), add_special_tokens=False)["input_ids"][:120] for f in fields]
    cls, sep = tok.cls_token_id, tok.sep_token_id
    out, cur, cur_idx = [], [], []
    budget = max_len - 2 - len(ctx_ids)
    used = 0
    for i, ids in enumerate(f_ids):
        if cur and used + len(ids) > budget:
            out.append((cur, cur_idx))
            cur, cur_idx, used = [], [], 0
        cur.append(ids)
        cur_idx.append(i)
        used += len(ids)
    if cur or not out:
        out.append((cur, cur_idx))
    res = []
    for chunk, idx in out:
        ids = [cls] + ctx_ids
        pos = []
        for f in chunk:
            pos.append(len(ids))  # f starts with the [FLD] token
            ids += f
        ids.append(sep)
        res.append((ids, pos, idx))
    return res


def augment(ctx, fields, rng: random.Random):
    """Hint dropout: teaches the model to cope with missing autocomplete/type/labels (e.g. revealed passwords)."""
    ctx = dict(ctx)
    for k, p in (("title", 0.1), ("url", 0.1), ("txt", 0.15), ("link", 0.1), ("hid", 0.2), ("head", 0.1)):
        if rng.random() < p:
            ctx[k] = ""
    out = []
    for f in fields:
        f = dict(f)
        if rng.random() < 0.25:
            f["ac"] = ""
        if f["type"] == "input password" and rng.random() < 0.08:
            f["type"] = "input text"
        for k, p in (("name", 0.05), ("id", 0.05), ("class", 0.1), ("lbl", 0.1), ("ph", 0.07), ("near", 0.1),
                     ("after", 0.1), ("data", 0.1)):
            if rng.random() < p:
                f[k] = ""
        out.append(f)
    return ctx, out


def train_tokenizer(texts_iter, out_dir, vocab_size=16384):
    import os
    from tokenizers import BertWordPieceTokenizer
    os.makedirs(out_dir, exist_ok=True)
    t = BertWordPieceTokenizer(lowercase=True, strip_accents=False, clean_text=True, handle_chinese_chars=True)
    t.train_from_iterator(texts_iter, vocab_size=vocab_size, min_frequency=3,
                          special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + SPECIAL_TOKENS,
                          limit_alphabet=6000, wordpieces_prefix="##")
    from transformers import PreTrainedTokenizerFast
    hf = PreTrainedTokenizerFast(tokenizer_object=t._tokenizer, unk_token="[UNK]", pad_token="[PAD]",
                                 cls_token="[CLS]", sep_token="[SEP]", mask_token="[MASK]",
                                 model_max_length=MAX_LEN)
    hf.save_pretrained(out_dir)
    return hf