"""Versioned standalone container for the independent codec's scalar indices.

Frame-major, stage-major, coordinate-major fixed-width symbols occupy exactly
46 bits/frame, packed MSB first. A fixed header stores model SHA-256, waveform
length and frame count. CRC32 covers the header fields and packed payload.
"""
from __future__ import annotations

from dataclasses import dataclass
import struct
import zlib

import numpy as np

MAGIC = b"TCSPCH1\0"
VERSION = 1
SAMPLE_RATE = 16000
HOP = 320
BITS_PER_FRAME = 46
LEVELS = np.array([[8, 8, 8, 4, 4, 4]] * 2 + [[4, 4, 2, 2, 2, 2]] * 2)
WIDTHS = np.log2(LEVELS).astype(np.int64).reshape(-1)
HEADER_WITHOUT_CRC = struct.Struct("<8sBBHIQQ32s")
CRC = struct.Struct("<I")
HEADER_BYTES = HEADER_WITHOUT_CRC.size + CRC.size
DEFAULT_MAX_FRAMES = 1000000


@dataclass(frozen=True)
class PackedSpeech:
    codes: np.ndarray
    samples: int
    model_sha256: str
    file_bytes: int
    sample_rate: int = SAMPLE_RATE

    @property
    def payload_bits(self):
        return self.codes.shape[-1] * BITS_PER_FRAME

    def accounting(self):
        seconds = self.samples / self.sample_rate
        payload_bytes = (self.payload_bits + 7) // 8
        return {"samples": self.samples, "frames": self.codes.shape[-1],
                "seconds": seconds, "sample_rate": self.sample_rate,
                "model_sha256": self.model_sha256, "payload_bits": self.payload_bits,
                "payload_bytes": payload_bytes,
                "padding_bits": payload_bytes * 8 - self.payload_bits,
                "header_bytes": HEADER_BYTES, "file_bytes": self.file_bytes,
                "payload_bits_per_second": self.payload_bits / seconds,
                "file_bits_per_second": self.file_bytes * 8 / seconds,
                "format": "tinycall-independent-packed-v1"}


def model_bytes(fingerprint):
    if not isinstance(fingerprint, str) or len(fingerprint) != 64:
        raise ValueError("Model identity must be a SHA-256 hex string")
    try:
        value = bytes.fromhex(fingerprint)
    except ValueError as error:
        raise ValueError("Model identity must be a SHA-256 hex string") from error
    if len(value) != 32:
        raise ValueError("Model SHA-256 must contain exactly 32 bytes")
    return value


def pack_codes(codes, samples, model_sha256):
    values = np.asarray(codes)
    if values.ndim != 3 or values.shape[:2] != (4, 6) or values.shape[-1] < 1:
        raise ValueError("Expected nonempty [4,6,frames] scalar indices")
    if values.dtype.kind not in "iu":
        raise ValueError("Scalar indices must have integer dtype")
    if np.any(values < 0) or np.any(values >= LEVELS[:, :, None]):
        raise ValueError("Scalar index outside its grid")
    if isinstance(samples, bool) or not isinstance(samples, (int, np.integer)) or samples <= 0:
        raise ValueError("Sample count must be a positive integer")
    frames = values.shape[-1]
    if frames > DEFAULT_MAX_FRAMES or frames != (int(samples) + HOP - 1) // HOP:
        raise ValueError("Frame count does not match waveform duration or exceeds format limit")
    words = np.zeros(frames, dtype=np.uint64)
    shift = BITS_PER_FRAME
    for coordinate, width in zip(values.reshape(24, frames), WIDTHS):
        shift -= int(width)
        words |= coordinate.astype(np.uint64) << np.uint64(shift)
    shifts = np.arange(BITS_PER_FRAME - 1, -1, -1, dtype=np.uint64)
    bits = ((words[:, None] >> shifts) & 1).astype(np.uint8).reshape(-1)
    payload = np.packbits(bits, bitorder="big").tobytes()
    header = HEADER_WITHOUT_CRC.pack(MAGIC, VERSION, 0, BITS_PER_FRAME,
                                     SAMPLE_RATE, int(samples), frames,
                                     model_bytes(model_sha256))
    checksum = zlib.crc32(payload, zlib.crc32(header))
    return header + CRC.pack(checksum) + payload


def unpack_codes(blob, expected_model_sha256=None, max_frames=DEFAULT_MAX_FRAMES):
    if len(blob) < HEADER_BYTES:
        raise ValueError("Truncated codec header")
    header = blob[:HEADER_WITHOUT_CRC.size]
    magic, version, flags, width, rate, samples, frames, fingerprint = HEADER_WITHOUT_CRC.unpack(header)
    if magic != MAGIC or version != VERSION or flags != 0:
        raise ValueError("Unsupported codec format, version or flags")
    if width != BITS_PER_FRAME or rate != SAMPLE_RATE:
        raise ValueError("Unsupported scalar layout or sample rate")
    if samples == 0 or frames == 0 or frames != (samples + HOP - 1) // HOP:
        raise ValueError("Invalid waveform length or frame count")
    if frames > max_frames or frames > DEFAULT_MAX_FRAMES:
        raise ValueError("Declared stream exceeds the frame limit")
    size = (frames * BITS_PER_FRAME + 7) // 8
    if len(blob) != HEADER_BYTES + size:
        raise ValueError("Truncated payload or unexpected trailing bytes")
    payload = blob[HEADER_BYTES:]
    expected_crc, = CRC.unpack(blob[HEADER_WITHOUT_CRC.size:HEADER_BYTES])
    if zlib.crc32(payload, zlib.crc32(header)) != expected_crc:
        raise ValueError("Codec stream checksum mismatch")
    if expected_model_sha256 is not None and fingerprint != model_bytes(expected_model_sha256):
        raise ValueError("Codec stream requires a different model checkpoint")
    padding = size * 8 - frames * BITS_PER_FRAME
    if padding and payload[-1] & ((1 << padding) - 1):
        raise ValueError("Nonzero unused padding bits")
    bits = np.unpackbits(np.frombuffer(payload, dtype=np.uint8), bitorder="big")
    bits = bits[:frames * BITS_PER_FRAME].reshape(frames, BITS_PER_FRAME)
    shifts = np.arange(BITS_PER_FRAME - 1, -1, -1, dtype=np.uint64)
    words = np.sum(bits.astype(np.uint64) << shifts, axis=1, dtype=np.uint64)
    values = np.empty((24, frames), dtype=np.int64)
    shift = BITS_PER_FRAME
    for index, width in enumerate(WIDTHS):
        shift -= int(width)
        values[index] = (words >> np.uint64(shift)) & np.uint64((1 << int(width)) - 1)
    return PackedSpeech(values.reshape(4, 6, frames), int(samples), fingerprint.hex(), len(blob))