"""
AmunPy — Byte-level data filters for AMUN snapshots.

This module implements the byte-level decoding filters specified by the
AMUN output format and referenced in snapshot XML metadata via the
`data_encoder` attribute.

The filters defined here operate strictly on raw byte streams and are
applied *after decompression* and *before NumPy dtype interpretation*.
They are independent of XML parsing, file I/O, and AMR logic.

Currently supported filters:
  - 'shuffle'    : restores per-element byte ordering
  - 'bytedelta'  : decodes byte-wise delta encoding

IMPORTANT NOTES:
- These implementations are behavior-preserving extractions of the
  original AmunXML internal decoding routines.
- The 'bytedelta' filter intentionally uses int8 wraparound arithmetic,
  matching the AMUN on-disk encoding semantics.
- The details below (signedness, accumulation order, reshape logic) are
  format-defined invariants. Any change must be validated against real
  AMUN snapshots and digest checks.

The purpose of isolating this module is to provide a single, testable,
and reusable implementation of AMUN byte filters, enabling safer
refactoring of higher-level snapshot parsing code.
"""
from __future__ import annotations

import numpy as np
from typing import Optional


class DataFilterError(ValueError):
    pass


def bytedelta_decode(a: bytes, dtype: str | np.dtype = "int64") -> bytes:
    s = np.dtype(dtype).itemsize
    n = len(a)
    if n % s != 0:
        raise DataFilterError(f"Bytedelta decode: input length {n} not divisible by itemsize {s} for dtype={dtype!r}")
    arr = np.frombuffer(a, dtype=np.int8).reshape((s, -1))
    dec = np.add.accumulate(arr, axis=1, dtype=np.int8)
    return dec.T.tobytes()


def shuffle_decode(a: bytes, dtype: str | np.dtype = "int64") -> bytes:
    s = np.dtype(dtype).itemsize
    n = len(a)
    if n % s != 0:
        raise DataFilterError(f"Shuffle decode: input length {n} not divisible by itemsize {s} for dtype={dtype!r}")
    arr = np.frombuffer(a, dtype=np.int8).reshape((s, -1))
    return arr.T.tobytes()


def filter_decode(
    data: bytes,
    *,
    data_encoder: Optional[str] = None,
    dtype: str | np.dtype = "int64",
) -> bytes:
    """
    Apply AMUN byte-level data filter specified by 'data_encoder'.

    - If data_encoder is None/missing/empty/'none': return data unchanged.
    - Supported: 'shuffle', 'bytedelta'
    - Raises DataFilterError for unsupported encoders.
    """
    if data_encoder is None:
        return data

    enc = str(data_encoder).strip().lower()
    if enc in ("", "none"):
        return data

    if enc == "shuffle":
        return shuffle_decode(data, dtype=dtype)

    if enc == "bytedelta":
        return bytedelta_decode(data, dtype=dtype)

    raise DataFilterError(f"Unsupported data_encoder={data_encoder!r}")