"""Spotlight VM as a `transformers` model. The model runs spotlight.c, built into build/ on first use."""
import ctypes, glob, json, os, subprocess

import torch
from huggingface_hub import snapshot_download
from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer


class SpotlightConfig(PretrainedConfig):
    model_type = "spotlight"


class SpotlightForCausalLM(PreTrainedModel):
    config_class = SpotlightConfig

    @classmethod
    def from_pretrained(cls, name_or_path, *args, **kwargs):
        folder = os.path.abspath(name_or_path if os.path.isdir(name_or_path) else snapshot_download(name_or_path))
        library = os.path.join(folder, "build", "libspotlight.so")
        if not os.path.exists(library):
            os.makedirs(os.path.dirname(library), exist_ok=True)
            source = os.path.join(folder, "spotlight.c")
            subprocess.run(["cc", "-O3", "-std=c99", "-shared", "-fPIC", "-DSPOTLIGHT_LIBRARY", source, "-lm", "-o", library], check=True)
        model = cls(SpotlightConfig.from_pretrained(folder))
        model.folder, model.lib = folder, ctypes.CDLL(library)
        model.lib.spotlight_generate.restype = ctypes.c_int64
        model.lib.spotlight_kept.restype = ctypes.POINTER(ctypes.c_int32)
        model.lib.spotlight_load(folder.encode())
        return model

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=10**12, packages=None, **kwargs):
        """Runs the program from a fresh memory: the core plus `packages` (default: all in memory/).
        Returns the prompt followed by the generated tokens that carry output, and the stop token."""
        if packages is None:
            packages = [p for p in sorted(glob.glob(os.path.join(self.folder, "memory", "*.safetensors")))
                        if not p.endswith("core.safetensors")]
        paths = (ctypes.c_char_p * max(1, len(packages)))(*[os.path.abspath(p).encode() for p in packages])
        prompt = input_ids[0].to(torch.int32).contiguous()
        self.generated_tokens = self.lib.spotlight_generate(ctypes.c_void_p(prompt.data_ptr()), ctypes.c_int64(prompt.numel()),
                                                            ctypes.c_int64(max_new_tokens), paths, len(packages))
        n = ctypes.c_int64()
        kept = ctypes.string_at(self.lib.spotlight_kept(ctypes.byref(n)), 4 * n.value)
        return torch.cat([prompt, torch.frombuffer(bytearray(kept), dtype=torch.int32)]).long()[None]


class SpotlightTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "intelligence/vocab.json", "config_file": "intelligence/config.json"}

    def __init__(self, vocab_file, config_file, **kwargs):
        with open(vocab_file) as f:
            self.names = json.load(f)
        self.ids = {name: i for i, name in enumerate(self.names)}
        with open(config_file) as f:
            self.input_base = json.load(f)["input_base"]
        super().__init__(**kwargs)

    vocab_size = property(lambda self: len(self.names))
    get_vocab = lambda self: dict(self.ids)
    _convert_token_to_id = lambda self, token: self.ids[token]
    _convert_id_to_token = lambda self, index: self.names[index]

    def _tokenize(self, text, **kwargs):
        """Where the source goes, its bytes (and a NUL) page by page, then execute -- as spotlight.c does."""
        data, base, tokens = text.encode("ascii") + b"\0", self.input_base, []
        for page in reversed(range((base % 65536 + len(data) - 1) // 65536 + 1)):
            address = (base // 65536 + page) * 65536 if page else base
            tokens += [f"{address >> 8 * lane & 255:02x}" for lane in range(4)] + ["page"]
        for i, b in enumerate(data):
            tokens += (["unpage"] if i and (base + i) % 65536 == 0 else []) + [chr(b) if 0x20 < b < 0x7F else f"ascii({b:02x})"]
        return tokens + ["execute"]

    def decode(self, token_ids, **kwargs):
        """The program's printed output: the bytes carried by the out(..) tokens."""
        names = [self.names[i] for i in (token_ids.tolist() if hasattr(token_ids, "tolist") else token_ids)]
        return bytes(ord(n[4]) if len(n) == 6 else int(n[4:-1], 16) for n in names if n.startswith("out(")).decode("latin-1")