"""Hashed n-gram addressing shared by Senryu-2's model, its host table store, and serving.

Each table row is addressed with integer ops that give the same result on CPU
and CUDA: canonical token ids (NFKC-normalized, lowercased, stripped surface
forms) are multiplied by odd 31-bit constants, XOR-combined, mixed, and reduced
modulo a distinct prime per table. Token ids stay below 2**16, so every
product stays below 2**47 and int64 never overflows.

Tables live in one concatenated ``[total_rows, embedding_dim]`` tensor; table
``t`` owns rows ``[offset_t, offset_t + rows_t)``. The prime row counts, the
multipliers, and the canonical map are saved beside every table file and
checked on load, so a table can never be read with a different address map.
"""

from __future__ import annotations

import unicodedata
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from tokenizers import Tokenizer

from tiny_gdn.config import TinyGDNConfig

NGRAM_TABLES_FILE = "ngram_tables.safetensors"
NGRAM_TABLES_EMA_FILE = "ngram_tables_ema.safetensors"
NGRAM_OPTIMIZER_FILE = "ngram_optimizer.safetensors"

PRIME_STRIDE = 4099
HASH_SEED = 0x6B696D6F
MIX_SHIFT = 29
LOAD_CHUNK_ROWS = 1 << 20
_MASK64 = (1 << 64) - 1
_MILLER_RABIN_BASES = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37)


def _splitmix64(state: int) -> tuple[int, int]:
    state = (state + 0x9E3779B97F4A7C15) & _MASK64
    mixed = ((state ^ (state >> 30)) * 0xBF58476D1CE4E5B9) & _MASK64
    mixed = ((mixed ^ (mixed >> 27)) * 0x94D049BB133111EB) & _MASK64
    return state, mixed ^ (mixed >> 31)


def _is_prime(value: int) -> bool:
    if value < 2:
        return False
    for base in _MILLER_RABIN_BASES:
        if value % base == 0:
            return value == base
    odd, twos = value - 1, 0
    while odd % 2 == 0:
        odd //= 2
        twos += 1
    for base in _MILLER_RABIN_BASES:
        witness = pow(base, odd, value)
        if witness in (1, value - 1):
            continue
        for _ in range(twos - 1):
            witness = witness * witness % value
            if witness == value - 1:
                break
        else:
            return False
    return True


def _next_prime(value: int) -> int:
    while not _is_prime(value):
        value += 1
    return value


@dataclass(frozen=True)
class NGramLayout:
    """Where every n-gram of every order lands in the concatenated tables."""

    orders: tuple[int, ...]
    heads_per_order: int
    table_rows: tuple[int, ...]
    table_offsets: tuple[int, ...]
    multipliers: tuple[tuple[int, ...], ...]
    embedding_dim: int
    pad_token_id: int

    @classmethod
    def from_config(cls, config: TinyGDNConfig) -> NGramLayout:
        if not config.ngram_orders:
            raise ValueError("config has no n-gram memory")
        rows: list[int] = []
        candidate = config.ngram_table_rows
        for _ in range(config.ngram_table_count):
            candidate = _next_prime(candidate + PRIME_STRIDE)
            rows.append(candidate)
        offsets = tuple(int(value) for value in np.cumsum([0, *rows[:-1]]))
        state = HASH_SEED
        multipliers: list[tuple[int, ...]] = []
        max_order = max(config.ngram_orders)
        for _ in range(config.ngram_table_count):
            table: list[int] = []
            for _ in range(max_order):
                state, value = _splitmix64(state)
                table.append((value >> 33) | (1 << 30) | 1)
            multipliers.append(tuple(table))
        return cls(
            orders=tuple(config.ngram_orders),
            heads_per_order=config.ngram_heads_per_order,
            table_rows=tuple(rows),
            table_offsets=offsets,
            multipliers=tuple(multipliers),
            embedding_dim=config.ngram_embedding_dim,
            pad_token_id=config.bos_token_id,
        )

    @property
    def table_count(self) -> int:
        return len(self.table_rows)

    @property
    def total_rows(self) -> int:
        return self.table_offsets[-1] + self.table_rows[-1]

    @property
    def max_order(self) -> int:
        return max(self.orders)

    @property
    def memory_dim(self) -> int:
        return self.table_count * self.embedding_dim

    @property
    def table_orders(self) -> tuple[int, ...]:
        return tuple(order for order in self.orders for _ in range(self.heads_per_order))

    def metadata_tensors(self) -> dict[str, torch.Tensor]:
        return {
            "table_rows": torch.tensor(self.table_rows, dtype=torch.int64),
            "hash_multipliers": torch.tensor(self.multipliers, dtype=torch.int64),
        }

    def check_metadata(self, tensors: dict[str, torch.Tensor], source: Path) -> None:
        for name, expected in self.metadata_tensors().items():
            found = tensors.get(name)
            if found is None or not torch.equal(found.cpu(), expected):
                raise ValueError(f"{source} was written for a different n-gram layout ({name})")


def canonical_token_ids(tokenizer: Tokenizer, vocab_size: int) -> torch.Tensor:
    """Map every token id to the smallest id sharing its normalized surface form.

    Normalization is NFKC, then strip, then lowercase. Special tokens keep
    their own ids, whitespace-only tokens are only NFKC-merged, and tokens
    that decode to partial UTF-8 stay distinct.
    """

    special = set(tokenizer.get_added_tokens_decoder())
    canonical = np.arange(vocab_size, dtype=np.int64)
    first_by_key: dict[tuple[str, str], int] = {}
    for token_id in range(min(vocab_size, tokenizer.get_vocab_size())):
        if token_id in special:
            continue
        text = tokenizer.decode([token_id], skip_special_tokens=False)
        if "\ufffd" in text:
            key = ("bytes", tokenizer.id_to_token(token_id) or str(token_id))
        else:
            normalized = unicodedata.normalize("NFKC", text)
            folded = normalized.strip().lower()
            key = ("text", folded) if folded else ("space", normalized)
        canonical[token_id] = first_by_key.setdefault(key, token_id)
    return torch.from_numpy(canonical)


def document_distance(tokens: torch.Tensor, bos_token_id: int, row_starts: torch.Tensor | None = None) -> torch.Tensor:
    """Offset of every token in a flat ``[N]`` stream from its document's first token.

    Documents start at position 0, at every BOS token, and at ``row_starts``.
    """

    starts = tokens == bos_token_id
    starts[0] = True
    if row_starts is not None:
        starts[row_starts] = True
    positions = torch.arange(tokens.numel(), device=tokens.device)
    first = torch.where(starts, positions, torch.zeros_like(positions)).cummax(dim=0).values
    return positions - first


def ngram_row_indices(
    canonical: torch.Tensor,
    layout: NGramLayout,
    *,
    distance: torch.Tensor | None = None,
    history: torch.Tensor | None = None,
) -> torch.Tensor:
    """Global table rows ``[..., L, tables]`` for the n-grams ending at each token.

    ``canonical`` holds int64 canonical ids ``[..., L]``. Tokens before a
    document start read as the pad id: with ``distance`` (same shape) a
    document starts wherever the distance is zero; with ``history`` (the
    ``max_order - 1`` canonical ids before position 0) the sequence continues
    a cached one; with neither, position 0 starts the only document.
    """

    span = layout.max_order
    length = canonical.shape[-1]
    pad = layout.pad_token_id
    if history is not None:
        if history.shape[-1] != span - 1:
            raise ValueError(f"history must hold {span - 1} tokens")
        extended = torch.cat([history, canonical], dim=-1)
        shifted = [extended[..., span - 1 - back : span - 1 - back + length] for back in range(span)]
    else:
        shifted = [canonical]
        for back in range(1, span):
            if back >= length:
                previous = torch.full_like(canonical, pad)
            else:
                padding = canonical.new_full((*canonical.shape[:-1], back), pad)
                previous = torch.cat([padding, canonical[..., :-back]], dim=-1)
            if distance is not None:
                previous = torch.where(distance >= back, previous, pad)
            shifted.append(previous)
    columns = []
    for table, order in enumerate(layout.table_orders):
        multipliers = layout.multipliers[table]
        mixed = shifted[0] * multipliers[0]
        for back in range(1, order):
            mixed = mixed ^ (shifted[back] * multipliers[back])
        mixed = mixed ^ (mixed >> MIX_SHIFT)
        columns.append(torch.remainder(mixed, layout.table_rows[table]) + layout.table_offsets[table])
    return torch.stack(columns, dim=-1)


def save_table_file(path: Path, tables: torch.Tensor, layout: NGramLayout, canonical: torch.Tensor) -> None:
    if tables.shape != (layout.total_rows, layout.embedding_dim):
        raise ValueError(f"tables have shape {tuple(tables.shape)}, layout needs {(layout.total_rows, layout.embedding_dim)}")
    save_file(
        {
            "tables": tables.contiguous(),
            "canonical_ids": canonical.to(dtype=torch.int64, device="cpu").contiguous(),
            **layout.metadata_tensors(),
        },
        str(path),
    )


def read_table_metadata(path: Path, layout: NGramLayout) -> torch.Tensor:
    """Check ``path`` against ``layout`` and return its canonical token map."""

    with safe_open(str(path), framework="pt", device="cpu") as handle:
        metadata = {name: handle.get_tensor(name) for name in ("table_rows", "hash_multipliers")}
        layout.check_metadata(metadata, path)
        # get_tensor returns a view of a mapping of the whole file, which a kept
        # view would hold open; FUSE mounts such as ntfs-3g then cannot delete it.
        return handle.get_tensor("canonical_ids").clone()


def copy_rows_from_file(path: Path, name: str, destination: torch.Tensor) -> None:
    """Copy tensor ``name`` of a safetensors file into ``destination`` in row chunks."""

    with safe_open(str(path), framework="pt", device="cpu") as handle:
        source = handle.get_slice(name)
        shape = tuple(source.get_shape())
        if shape != tuple(destination.shape):
            raise ValueError(f"{path} holds {name} of shape {shape}, expected {tuple(destination.shape)}")
        for start in range(0, shape[0], LOAD_CHUNK_ROWS):
            stop = min(start + LOAD_CHUNK_ROWS, shape[0])
            destination[start:stop].copy_(source[start:stop])


def load_table_file(
    path: Path,
    layout: NGramLayout,
    *,
    device: torch.device | str = "cpu",
    dtype: torch.dtype = torch.float32,
    out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """``(tables, canonical_ids)`` from ``path``, copied in row chunks to bound host memory."""

    canonical = read_table_metadata(path, layout)
    tables = (
        torch.empty((layout.total_rows, layout.embedding_dim), dtype=dtype, device=device)
        if out is None
        else out
    )
    copy_rows_from_file(path, "tables", tables)
    return tables, canonical