"""Run termite-classifier with PyTorch. It reads plain data files and never runs a command.
Usage: python termite_classifier.py 'ls -la' [FOLDER] [SHELL] [OS] [SESSION]"""
import json, pathlib, struct, sys
import torch, torch.nn as nn, torch.nn.functional as F

HERE = pathlib.Path(__file__).resolve().parent
QUESTIONS = ("only_reads", "no_secrets", "harmful")
TYPES = {"F32": torch.float32}
CLS, FIRST = 1, 3


def pieces(data):
    """A piece is a run of ASCII letters, digits or bytes over 127, or one other byte."""
    start = 0
    for i, b in enumerate(data):
        if not (48 <= b <= 57 or 65 <= b <= 90 or 97 <= b <= 122 or b > 127):
            if i > start:
                yield data[start:i]
            yield data[i:i + 1]
            start = i + 1
    if start < len(data):
        yield data[start:]


class Tokeniser:
    def __init__(self, merges):
        self.rank = {tuple(m): i for i, m in enumerate(merges)}

    def piece(self, data):
        w = [FIRST + b for b in data]
        while len(w) > 1:
            found = min((self.rank.get(p, 1 << 30), i) for i, p in enumerate(zip(w, w[1:])))
            if found[0] == 1 << 30:
                break
            w[found[1]:found[1] + 2] = [FIRST + 256 + found[0]]
        return w

    def encode(self, text):
        ids = [CLS]
        for p in pieces(text.encode("utf8", "surrogateescape")):
            ids += self.piece(p)
        return ids


class Block(nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        self.heads = heads
        self.n1, self.n2 = nn.LayerNorm(width), nn.LayerNorm(width)
        self.qkv, self.out = nn.Linear(width, 3 * width), nn.Linear(width, width)
        self.up, self.down = nn.Linear(width, 4 * width), nn.Linear(4 * width, width)

    def forward(self, x, keep):
        b, t, w = x.shape
        q, k, v = self.qkv(self.n1(x)).view(b, t, 3, self.heads, w // self.heads).permute(2, 0, 3, 1, 4)
        a = F.scaled_dot_product_attention(q, k, v, attn_mask=keep[:, None, None, :])
        x = x + self.out(a.transpose(1, 2).reshape(b, t, w))
        return x + self.down(F.gelu(self.up(self.n2(x))))


class Encoder(nn.Module):
    def __init__(self, vocab, longest, layers, width, heads):
        super().__init__()
        self.tokens, self.places = nn.Embedding(vocab, width), nn.Embedding(longest, width)
        self.blocks = nn.ModuleList(Block(width, heads) for _ in range(layers))
        self.last = nn.LayerNorm(width)
        self.heads = nn.Linear(width, 3)

    def forward(self, ids):
        keep = ids != 0
        x = self.tokens(ids) + self.places(torch.arange(ids.shape[1]))
        for block in self.blocks:
            x = block(x, keep)
        return self.heads(self.last(x)[:, 0])


def read_weights(path):
    with open(path, "rb") as f:
        size = struct.unpack("<Q", f.read(8))[0]
        header = json.loads(f.read(size))
        body = bytearray(f.read())
    return {name: torch.frombuffer(body, dtype=TYPES[t["dtype"]], offset=t["data_offsets"][0],
                                   count=(t["data_offsets"][1] - t["data_offsets"][0]) // 4).view(t["shape"])
            for name, t in header.items() if name != "__metadata__"}


class Classifier:
    def __init__(self, folder=HERE):
        folder = pathlib.Path(folder)
        self.config = json.load(open(folder / "config.json"))
        c = self.config
        self.tok = Tokeniser(json.load(open(folder / "merges.json")))
        self.net = Encoder(c["vocab"], c["longest"], c["layers"], c["width"], c["heads"])
        self.net.load_state_dict(read_weights(folder / c["weights"]))
        self.net.eval()

    @torch.no_grad()
    def answer(self, text):
        """Chances for each question. A command too long to read gets the strict answer."""
        ids = self.tok.encode(text)
        if len(ids) > self.config["longest"]:
            return {"only_reads": 0.0, "no_secrets": 0.0, "harmful": 1.0}
        z = self.net(torch.tensor([ids]))[0]
        return {q: torch.sigmoid(z[i] / self.config["temperatures"][i]).item() for i, q in enumerate(QUESTIONS)}

    def safe(self, a):
        limit = self.config["limits"]
        return a["only_reads"] >= limit["only_reads"] and a["no_secrets"] >= limit["no_secrets"] and a["harmful"] <= limit["harmful"]


if __name__ == "__main__":
    command, *rest = sys.argv[1:]
    folder, shell, system, session = (rest + ["/home/dev/src/project", "bash", "linux", "local"][len(rest):])[:4]
    model = Classifier()
    answer = model.answer(f"{shell} {system} {session}\n{folder}\n{command}")
    print(json.dumps({**answer, "safe": model.safe(answer)}, indent=1))