"""Prime dataset: n -> next prime, digit-tokenized; splits and batching."""
import random

import torch

from src.config import Config


def sieve_primes(limit: int) -> list[int]:
    """All primes <= limit (inclusive)."""
    if limit < 2:
        return []
    is_prime = [True] * (limit + 1)
    is_prime[0] = is_prime[1] = False
    for p in range(2, int(limit ** 0.5) + 1):
        if is_prime[p]:
            for m in range(p * p, limit + 1, p):
                is_prime[m] = False
    return [i for i in range(2, limit + 1) if is_prime[i]]


def next_prime(n: int, primes: list[int]) -> int:
    for p in primes:
        if p > n:
            return p
    raise ValueError(f"no prime > {n} in supplied list")


def encode_int(n: int, cfg: Config) -> list[int]:
    if cfg.vocab_mode == "integers":
        return [n]
    return [int(d) for d in str(n)]


def decode_tokens(ts, cfg: Config) -> int:
    """Decode a token sequence, stopping at EOS. -1 if nothing decodable."""
    if cfg.vocab_mode == "integers":
        return int(ts[0]) if len(ts) else -1
    digits = []
    for t in ts:
        if t == cfg.eos_id:
            break
        digits.append(int(t))
    return int("".join(map(str, digits))) if digits else -1


def get_splits(cfg: Config) -> tuple[list[int], list[int]]:
    """(train, val) input lists, seeded shuffle, no overlap."""
    rng = random.Random(cfg.seed)
    inputs = list(range(cfg.range_start, cfg.range_end + 1))
    rng.shuffle(inputs)
    n_val = max(1, round(len(inputs) * cfg.holdout_frac))
    return sorted(inputs[n_val:]), sorted(inputs[:n_val])


def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list[int]]]:
    """[(input_tokens, target_tokens+EOS), ...]"""
    primes = sieve_primes(cfg.range_end + 100)
    out = []
    for n in inputs:
        p = next_prime(n, primes)
        out.append((encode_int(n, cfg), encode_int(p, cfg) + [cfg.eos_id]))
    return out


def _global_lengths(cfg: Config) -> tuple[int, int]:
    """(in_max, out_max): fixed global lengths so batch layout == singleton layout (codex BLOCKER fix)."""
    primes = sieve_primes(cfg.range_end + 100)
    max_target = next_prime(cfg.range_end, primes)
    if cfg.vocab_mode == "integers":
        return 1, 2                      # [value], [value, EOS]
    return len(str(cfg.range_end)), len(str(max_target)) + 1   # digits + EOS


def pad_inputs(x: torch.Tensor, cfg: Config) -> torch.Tensor:
    """LEFT-pad inputs to the global in_max so absolute positions are layout-invariant."""
    in_max, _ = _global_lengths(cfg)
    if x.shape[1] < in_max:
        pad = torch.full((x.shape[0], in_max - x.shape[1]), cfg.pad_id, dtype=x.dtype, device=x.device)
        x = torch.cat([pad, x], dim=1)
    return x


def make_batch(examples, cfg: Config) -> dict[str, torch.Tensor]:
    """Fixed global layout (not batch-max): x left-padded to in_max, y right-padded to out_max."""
    xs, ys = zip(*examples)
    in_max, out_max = _global_lengths(cfg)
    B = len(examples)
    x = torch.full((B, in_max), cfg.pad_id, dtype=torch.long)
    y = torch.full((B, out_max), cfg.pad_id, dtype=torch.long)
    for i, (xi, yi) in enumerate(examples):
        x[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long)
        y[i, : len(yi)] = torch.tensor(yi, dtype=torch.long)
    # teacher-forced decoder input: BOS(=EOS reuse) then shifted y
    y_in = torch.cat([torch.full((B, 1), cfg.eos_id, dtype=torch.long), y[:, :-1]], dim=1)
    y_mask = y != cfg.pad_id
    return {"x": x, "y": y, "y_in": y_in, "y_mask": y_mask}