"""Stdlib-only WAV helpers (no numpy). SDK audio is interleaved float32 in
[-1,1]. read_wav accepts 16-bit PCM or 32-bit IEEE float input; save_wav
writes 16-bit PCM."""

import array
import struct
import sys
import wave

_INT16_MAX = 32767
# 2^15: decode divisor so -32768 maps to -1.0 and all samples stay within [-1, 1].
_INT16_SCALE = 32768.0

# WAV format tags (wFormatTag in the fmt chunk).
_WAVE_FORMAT_PCM = 0x0001
_WAVE_FORMAT_IEEE_FLOAT = 0x0003
_WAVE_FORMAT_EXTENSIBLE = 0xFFFE


def read_wav(path):
    """Read a 16-bit PCM or 32-bit IEEE float WAV.

    Returns (float32_bytes, sample_rate, channels) where float32_bytes is
    interleaved little-endian float32 in [-1, 1].

    Python's ``wave`` module only supports PCM, so the header is parsed
    directly to also accept IEEE-float (and WAVE_FORMAT_EXTENSIBLE) files.
    """
    with open(path, "rb") as f:
        riff, _size, wave_id = struct.unpack("<4sI4s", f.read(12))
        if riff != b"RIFF" or wave_id != b"WAVE":
            raise ValueError(f"{path}: not a RIFF/WAVE file")

        fmt_tag = channels = sample_rate = bits = None
        data = None

        while True:
            header = f.read(8)
            if len(header) < 8:
                break
            chunk_id, chunk_size = struct.unpack("<4sI", header)
            if chunk_id == b"fmt ":
                fmt = f.read(chunk_size)
                (fmt_tag, channels, sample_rate, _byte_rate, _block_align, bits) = (
                    struct.unpack("<HHIIHH", fmt[:16])
                )
                # WAVE_FORMAT_EXTENSIBLE stores the real tag in the SubFormat
                # GUID's first two bytes.
                if fmt_tag == _WAVE_FORMAT_EXTENSIBLE and len(fmt) >= 26:
                    fmt_tag = struct.unpack("<H", fmt[24:26])[0]
            elif chunk_id == b"data":
                data = f.read(chunk_size)
            else:
                f.seek(chunk_size, 1)  # skip unrelated chunk (LIST, fact, ...)
            # Chunks are word-aligned: skip the pad byte after an odd size.
            if chunk_size % 2:
                f.seek(1, 1)

    if fmt_tag is None or data is None:
        raise ValueError(f"{path}: missing fmt or data chunk")

    if fmt_tag == _WAVE_FORMAT_PCM and bits == 16:
        ints = array.array("h")
        ints.frombytes(data)
        if sys.byteorder == "big":
            ints.byteswap()  # WAV data is little-endian
        floats = array.array("f", (s / _INT16_SCALE for s in ints))
    elif fmt_tag == _WAVE_FORMAT_IEEE_FLOAT and bits == 32:
        floats = array.array("f")
        floats.frombytes(data)
        if sys.byteorder == "big":
            floats.byteswap()
    else:
        raise ValueError(
            f"{path}: unsupported format (tag={fmt_tag}, bits={bits}); "
            "expected 16-bit PCM or 32-bit IEEE float"
        )

    # Emit little-endian float32 bytes regardless of host endianness.
    if sys.byteorder == "big":
        out = array.array("f", floats)
        out.byteswap()
        return out.tobytes(), sample_rate, channels
    return floats.tobytes(), sample_rate, channels


def save_wav(path, float32_bytes, sample_rate, channels):
    """Write interleaved float32 bytes -> 16-bit PCM WAV."""
    floats = array.array("f")
    floats.frombytes(float32_bytes)
    ints = array.array(
        "h",
        (int(max(-1.0, min(1.0, f)) * _INT16_MAX) for f in floats),
    )
    with wave.open(path, "wb") as w:
        w.setnchannels(channels)
        w.setsampwidth(2)
        w.setframerate(sample_rate)
        w.writeframes(ints.tobytes())
