"""无损 WAV 读写核心。

设计要点（面向"高质量、高清晰度"）：

* 支持 16 / 24 / 32 位整数 PCM 与 32 位浮点（IEEE float）四种载荷；
* 24 位采用真正的 3 字节打包，不做 16 位截断；
* 多声道、24 位、浮点载荷统一使用 ``WAVE_FORMAT_EXTENSIBLE``（含通道掩码），
  避免播放器把 24 位误判为 16 位、把浮点误判为整数；
* 高位深写入时可选 **TPDF（三角概率密度）抖动**，把量化失真从
  "与信号相关的非线性失真" 变成 "与信号无关的宽带白噪"，这是后期降位深的
  标准做法，听感明显更干净；
* 录音过程中定期回写文件头（``checkpoint``），即使程序崩溃/断电，
  已落盘的数据也能被播放器正常识别；
* 超过 4 GiB 时可选 RF64（EBU Tech 3306）容器，避免 RIFF 32 位长度溢出。

本模块不依赖 numpy 以外的任何第三方库。
"""

from __future__ import annotations

import os
import struct
from dataclasses import dataclass
from typing import BinaryIO, Iterator

import numpy as np

WAVE_FORMAT_PCM = 0x0001
WAVE_FORMAT_IEEE_FLOAT = 0x0003
WAVE_FORMAT_EXTENSIBLE = 0xFFFE

# KSDATAFORMAT_SUBTYPE_*
_SUBTYPE_PCM = bytes.fromhex("0100000000001000800000aa00389b71")
_SUBTYPE_FLOAT = bytes.fromhex("0300000000001000800000aa00389b71")

# 声道掩码（dwChannelMask），用于 EXTENSIBLE 头部
_CHANNEL_MASKS = {
    1: 0x4,        # SPEAKER_FRONT_CENTER -> 单声道习惯上用 CENTER，但很多软件写 FRONT_LEFT
    2: 0x3,        # FRONT_LEFT | FRONT_RIGHT
    4: 0x33,       # FL FR BL BR
    6: 0x3F,
    8: 0x63F,
}

# 1 GiB 保留量：RIFF 数据块上限 4 GiB-1，留出余量给头部与安全边界
RIFF_DATA_LIMIT = 0xFFF00000 - 64

__all__ = [
    "WavFormat",
    "WavWriter",
    "WavReader",
    "wav_info",
    "read_wav",
    "write_wav",
    "RIFF_DATA_LIMIT",
    "SUPPORTED_BIT_DEPTHS",
]

# (名称, 位深, 采样格式) —— 直接对应界面上的"位深"下拉框
SUPPORTED_BIT_DEPTHS = ("16", "24", "32", "float32")


@dataclass(frozen=True)
class WavFormat:
    """描述一个 WAV 载荷格式。"""

    samplerate: int
    channels: int
    encoding: str  # 'pcm' | 'float'
    bits: int      # 每个采样的有效位宽（float 为 32）

    @property
    def bytes_per_sample(self) -> int:
        return self.bits // 8

    @property
    def frame_bytes(self) -> int:
        return self.bytes_per_sample * self.channels

    @property
    def bytes_per_second(self) -> float:
        return float(self.frame_bytes) * self.samplerate

    @property
    def block_align(self) -> int:
        return self.frame_bytes

    @property
    def label(self) -> str:
        if self.encoding == "float":
            return "32-bit float"
        return f"{self.bits}-bit PCM"

    def describe(self) -> str:
        return (
            f"{self.samplerate} Hz / {self.label} / "
            f"{'单声道' if self.channels == 1 else f'{self.channels} 声道'}"
        )


def _fmt_for_bitdepth(bit_depth: int | str) -> tuple[str, int]:
    """把界面上的位深选项翻译成 (encoding, bits)。"""
    if isinstance(bit_depth, str):
        s = bit_depth.strip().lower()
        if s in ("float", "float32", "f32", "32f"):
            return "float", 32
        s = s.rstrip("bit").rstrip("-").strip()
        bit_depth = int(s)
    bit_depth = int(bit_depth)
    if bit_depth == 16:
        return "pcm", 16
    if bit_depth == 24:
        return "pcm", 24
    if bit_depth == 32:
        return "pcm", 32
    if bit_depth == 64:
        raise ValueError("不支持 64 位整数 PCM；如需更高精度请使用 32-bit float")
    raise ValueError(f"不支持的位深: {bit_depth}")


def _safe_scale(bits: int) -> float:
    """整数满量程对应的浮点值。

    采用 ``2**(bits-1) - 1``（对称满量程），这样 +1.0 与 -1.0 都能被精确表示，
    不会出现 "正半周先削顶、负半周仍有余量" 的非对称削波。
    """
    return float((1 << (bits - 1)) - 1)


class _TpdfDither:
    """TPDF 抖动噪声源。

    三角分布 = 两个独立均匀分布之和，峰峰值恰好 1 LSB，
    这是业界公认的"无调制噪声"抖动，优于矩形（RPDF）抖动。
    """

    def __init__(self, seed: int | None = None):
        self._rng = np.random.Generator(np.random.PCG64(seed))

    def generate(self, size: int | tuple[int, ...], lsb: float = 1.0) -> np.ndarray:
        u1 = self._rng.random(size)
        u2 = self._rng.random(size)
        return (u1 - u2) * lsb

    def reseed(self, seed: int) -> None:
        self._rng = np.random.Generator(np.random.PCG64(seed))


def float_to_pcm(data: np.ndarray, encoding: str, bits: int,
                 dither: _TpdfDither | None = None) -> bytes:
    """把 float32/float64 的 [-1, 1] 数据编码为 WAV 载荷字节。"""
    x = np.asarray(data)
    if encoding == "float":
        return np.ascontiguousarray(x, dtype="<f4").tobytes()

    # 抖动与量化都必须在"整数域"内进行：抖动幅度 = 1 LSB（峰峰值）
    scale = _safe_scale(bits)
    y = x.astype(np.float64, copy=True) * scale
    if dither is not None:
        y += dither.generate(y.shape, lsb=1.0)

    ints = np.rint(y)
    lo, hi = -(1 << (bits - 1)), (1 << (bits - 1)) - 1
    np.clip(ints, lo, hi, out=ints)
    ints = ints.astype(np.int64)

    if bits == 16:
        return ints.astype("<i2").tobytes()
    if bits == 32:
        return ints.astype("<i4").tobytes()
    if bits == 24:
        buf = np.ascontiguousarray(ints, dtype="<i4")
        # 小端序下低 3 字节即为 24 位补码，直接切片比逐样本运算快两个数量级
        return buf.view(np.uint8).reshape(-1, 4)[:, :3].tobytes()
    raise ValueError(f"不支持的位深: {bits}")


def pcm_to_float(raw: bytes, encoding: str, bits: int) -> np.ndarray:
    """WAV 载荷字节 -> float64 数组（[-1, 1] 标度）。"""
    if encoding == "float":
        return np.frombuffer(raw, dtype="<f4").astype(np.float64)
    if bits == 16:
        return np.frombuffer(raw, dtype="<i2").astype(np.float64) / _safe_scale(16)
    if bits == 32:
        return np.frombuffer(raw, dtype="<i4").astype(np.float64) / _safe_scale(32)
    if bits == 24:
        b = np.frombuffer(raw, dtype=np.uint8)
        usable = (b.size // 3) * 3
        b = b[:usable].reshape(-1, 3).astype(np.int32)
        v = b[:, 0] | (b[:, 1] << 8) | (b[:, 2] << 16)
        # 手动符号扩展
        v = np.where(v & 0x800000, v - 0x1000000, v)
        return v.astype(np.float64) / _safe_scale(24)
    raise ValueError(f"不支持的位深: {bits}")


class WavWriter:
    """流式（边录边写）WAV 写入器。

    典型用法::

        w = WavWriter("a.wav", WavFormat(48000, 2, "pcm", 24), dither=True)
        w.write(block)          # block: float32 (n, channels)
        w.checkpoint()          # 每秒调用一次，保证崩溃安全
        info = w.close()        # 回写文件头，返回统计信息
    """

    def __init__(self, path: str, fmt: WavFormat, *, dither: bool = True,
                 rf64: bool = False, seed: int | None = None,
                 checkpoint_interval: float = 1.0):
        self.path = os.fspath(path)
        self.fmt = fmt
        self.rf64 = bool(rf64)
        self._dither = _TpdfDither(seed) if (dither and fmt.encoding == "pcm") else None
        self._checkpoint_interval = float(checkpoint_interval)
        self._last_checkpoint = 0.0

        self.frames_written = 0
        self.bytes_written = 0
        self.clipped_samples = 0
        self.peak = 0.0
        self.closed = False

        self._extensible = (
            fmt.encoding == "float" or fmt.bits > 16 or fmt.channels > 2
        )
        self._fh: BinaryIO = open(self.path, "wb", buffering=1024 * 1024)
        hdr = self._build_header()
        self._header = hdr
        self._data_size_offset = hdr.data_offset - 4
        self._fh.write(hdr.bytes)
        self._fh.flush()

    # ---------------------------------------------------------------- header
    def _build_header(self) -> "_Header":
        f = self.fmt
        if f.encoding == "float":
            fmt_tag = WAVE_FORMAT_IEEE_FLOAT if not self._extensible else WAVE_FORMAT_EXTENSIBLE
            subtype = _SUBTYPE_FLOAT
        else:
            fmt_tag = WAVE_FORMAT_PCM if not self._extensible else WAVE_FORMAT_EXTENSIBLE
            subtype = _SUBTYPE_PCM

        if self._extensible:
            mask = _CHANNEL_MASKS.get(f.channels, 0)
            fmt_body = struct.pack(
                "<HHIIHHHHI16s",
                WAVE_FORMAT_EXTENSIBLE, f.channels, f.samplerate,
                int(f.bytes_per_second), f.block_align, f.bits,
                22, f.bits, mask, subtype,
            )
        else:
            fmt_body = struct.pack(
                "<HHIIHH",
                fmt_tag, f.channels, f.samplerate,
                int(f.bytes_per_second), f.block_align, f.bits,
            )

        chunks: list[bytes] = []
        # fmt
        chunks.append(b"fmt " + struct.pack("<I", len(fmt_body)) + fmt_body)
        # fact（非 PCM 载荷必须声明采样帧数）
        if f.encoding == "float":
            chunks.append(b"fact" + struct.pack("<II", 4, 0))
            self._fact_offset = 12 + len(chunks[0]) + 8
        else:
            self._fact_offset = None

        riff_overhead = 12 + sum(len(c) for c in chunks) + 8
        if self.rf64:
            # RF64: ds64 紧跟 WAVE 之后，RIFF size 写 0xFFFFFFFF
            ds64_body = struct.pack("<QQQI", 0, 0, 0, 0)
            ds64 = b"ds64" + struct.pack("<I", len(ds64_body)) + ds64_body
            body = ds64 + b"".join(chunks)
            header = (b"RF64" + struct.pack("<I", 0xFFFFFFFF) + b"WAVE"
                      + body + b"data" + struct.pack("<I", 0))
            return _Header(header, len(header), riff_overhead, ds64_offset=12)
        body = b"".join(chunks)
        size = len(body) + 8 + 8  # 占位，收尾时回写
        header = (b"RIFF" + struct.pack("<I", size) + b"WAVE" + body
                  + b"data" + struct.pack("<I", 0))
        return _Header(header, len(header), riff_overhead, ds64_offset=None)

    # ----------------------------------------------------------------- write
    def write(self, data: np.ndarray) -> int:
        """写入一块浮点音频（形状 (n,) 或 (n, channels)），返回写入的采样帧数。"""
        if self.closed:
            raise ValueError("写入器已关闭")
        x = np.asarray(data, dtype=np.float32)
        if x.ndim == 1:
            x = x.reshape(-1, 1)
        if x.shape[1] != self.fmt.channels:
            raise ValueError(
                f"通道数不匹配：收到 {x.shape[1]}，期望 {self.fmt.channels}")

        n = int(x.shape[0])
        if n == 0:
            return 0

        # 统计（在抖动/削波之前，反映真实输入）
        blk_peak = float(np.max(np.abs(x))) if n else 0.0
        if blk_peak > self.peak:
            self.peak = blk_peak
        if blk_peak > 1.0:
            self.clipped_samples += int(np.count_nonzero(np.abs(x) > 1.0))

        payload = float_to_pcm(x, self.fmt.encoding, self.fmt.bits, self._dither)

        # 先判断是否越界，再落盘：这样调用方可以无损地切到下一个文件
        if not self.rf64 and self.bytes_written + len(payload) > RIFF_DATA_LIMIT:
            raise _RiffOverflow(
                f"RIFF 容器已达 4 GiB 上限（已写 {self.bytes_written} 字节），"
                "请启用自动分段或 RF64 模式")

        self._fh.write(payload)
        self.frames_written += n
        self.bytes_written += len(payload)
        return n

    def checkpoint(self, now: float | None = None) -> None:
        """定期回写长度字段，保证异常退出后文件仍可播放。"""
        if self.closed:
            return
        now = _monotonic() if now is None else now
        if now - self._last_checkpoint < self._checkpoint_interval:
            return
        self._last_checkpoint = now
        self._patch_sizes(final=False)

    def _patch_sizes(self, *, final: bool) -> None:
        f = self.fmt
        data_bytes = self.frames_written * f.frame_bytes
        pos = self._fh.tell()
        if self.rf64:
            body_len = len(self._header.bytes) - 8 + data_bytes
            if final:
                self._fh.seek(4)
                self._fh.write(struct.pack("<I", 0xFFFFFFFF))
                self._fh.seek(self._header.ds64_offset + 8)  # ds64 body 起点
                self._fh.write(struct.pack(
                    "<QQQI", body_len, data_bytes, self.frames_written, 0))
            self._fh.seek(pos)
            return
        # RIFF
        riff_size = (len(self._header.bytes) - 8) + data_bytes
        self._fh.seek(4)
        self._fh.write(struct.pack("<I", riff_size & 0xFFFFFFFF))
        self._fh.seek(self._data_size_offset)
        self._fh.write(struct.pack("<I", data_bytes & 0xFFFFFFFF))
        if self._fact_offset is not None:
            self._fh.seek(self._fact_offset)
            self._fh.write(struct.pack("<I", self.frames_written & 0xFFFFFFFF))
        self._fh.seek(pos)

    # ----------------------------------------------------------------- close
    def close(self) -> dict:
        if self.closed:
            return self.stats()
        self._patch_sizes(final=True)
        self._fh.flush()
        try:
            os.fsync(self._fh.fileno())
        except OSError:
            pass
        self._fh.close()
        self.closed = True
        return self.stats()

    def stats(self) -> dict:
        f = self.fmt
        duration = self.frames_written / f.samplerate if f.samplerate else 0.0
        return {
            "path": self.path,
            "frames": self.frames_written,
            "duration": duration,
            "bytes": self.bytes_written,
            "peak": self.peak,
            "peak_dbfs": 20 * np.log10(self.peak) if self.peak > 0 else float("-inf"),
            "clipped_samples": self.clipped_samples,
            "format": f.describe(),
            "dither": self._dither is not None,
            "rf64": self.rf64,
        }

    def abort(self) -> None:
        """放弃写入：关闭句柄但不回写头部（用于创建后立即失败的场景）。"""
        if not self.closed:
            try:
                self._fh.close()
            finally:
                self.closed = True

    def __enter__(self) -> "WavWriter":
        return self

    def __exit__(self, *exc) -> None:
        if exc and exc[0] is not None:
            self.abort()
        else:
            self.close()


@dataclass
class _Header:
    bytes: bytes
    data_offset: int
    riff_overhead: int
    ds64_offset: int | None = None


class _RiffOverflow(Exception):
    """RIFF 4 GiB 上限。引擎捕获后自动分段。"""


def _monotonic() -> float:
    import time
    return time.monotonic()


class WavReader:
    """轻量 WAV 读取器，支持 PCM 16/24/32、float32、EXTENSIBLE、RF64。"""

    def __init__(self, path: str):
        self.path = os.fspath(path)
        self.fh = open(self.path, "rb")
        self._parse()

    def _parse(self) -> None:
        fh = self.fh
        hdr = fh.read(12)
        if len(hdr) < 12:
            raise ValueError("不是有效的 WAV 文件（文件过短）")
        riff, _size, wave = hdr[:4], hdr[4:8], hdr[8:12]
        if riff == b"RF64":
            self.rf64 = True
        elif riff == b"RIFF":
            self.rf64 = False
        else:
            raise ValueError("不是有效的 WAV 文件（缺少 RIFF/RF64 标记）")
        if wave != b"WAVE":
            raise ValueError("不是有效的 WAV 文件（缺少 WAVE 标记）")

        self.fmt_tag = None
        self.channels = self.samplerate = self.bits = 0
        self.encoding = "pcm"
        self.data_offset = self.data_size = 0
        self.frames = 0
        self._ds64_data_size = None
        self._ds64_frames = None
        self._fact_frames = None

        while True:
            cid = fh.read(4)
            if len(cid) < 4:
                break
            (csize,) = struct.unpack("<I", fh.read(4))
            body_start = fh.tell()
            if cid == b"fmt ":
                body = fh.read(csize)
                (tag, ch, sr, _bps, _align, bits) = struct.unpack("<HHIIHH", body[:16])
                if tag == WAVE_FORMAT_EXTENSIBLE and csize >= 40:
                    (bits,) = struct.unpack("<H", body[18:20])
                    if body[24:40] == _SUBTYPE_FLOAT:
                        tag = WAVE_FORMAT_IEEE_FLOAT
                    elif body[24:40] == _SUBTYPE_PCM:
                        tag = WAVE_FORMAT_PCM
                self.fmt_tag, self.channels, self.samplerate = tag, ch, sr
                self.bits = bits
                self.encoding = "float" if tag == WAVE_FORMAT_IEEE_FLOAT else "pcm"
            elif cid == b"data":
                self.data_offset = body_start
                self.data_size = csize
            elif cid == b"fact" and csize >= 4:
                (self._fact_frames,) = struct.unpack("<I", fh.read(4))
            elif cid == b"ds64":
                body = fh.read(csize)
                _riff64, data64, frames64, _table = struct.unpack("<QQQI", body[:28])
                self._ds64_data_size = data64
                self._ds64_frames = frames64
            fh.seek(body_start + csize + (csize & 1))

        if self.fmt_tag is None:
            raise ValueError("WAV 文件缺少 fmt 块")
        if self.encoding == "pcm" and self.bits not in (16, 24, 32):
            raise ValueError(f"暂不支持的 PCM 位深: {self.bits}")
        if self.rf64 and self._ds64_data_size is not None:
            self.data_size = int(self._ds64_data_size)
        if self.channels <= 0 or self.samplerate <= 0:
            raise ValueError("WAV 头部信息不完整")
        self.frames = self.data_size // (self.channels * (self.bits // 8))

    @property
    def fmt(self) -> WavFormat:
        return WavFormat(self.samplerate, self.channels, self.encoding, self.bits)

    @property
    def duration(self) -> float:
        return self.frames / self.samplerate if self.samplerate else 0.0

    def read(self, start_frame: int = 0, num_frames: int | None = None) -> np.ndarray:
        """读取为 float64 数组，形状 (frames, channels)，范围 [-1, 1]。"""
        bps = self.bits // 8
        if num_frames is None:
            num_frames = max(0, self.frames - start_frame)
        num_frames = min(num_frames, max(0, self.frames - start_frame))
        if num_frames <= 0:
            return np.zeros((0, self.channels), dtype=np.float64)
        self.fh.seek(self.data_offset + start_frame * self.channels * bps)
        raw = self.fh.read(num_frames * self.channels * bps)
        flat = pcm_to_float(raw, self.encoding, self.bits)
        usable = (flat.size // self.channels) * self.channels
        return flat[:usable].reshape(-1, self.channels)

    def iter_blocks(self, block_frames: int = 1 << 16) -> Iterator[np.ndarray]:
        """分块迭代读取，便于处理超长录音而不占内存。"""
        pos = 0
        while pos < self.frames:
            block = self.read(pos, block_frames)
            if block.size == 0:
                break
            yield block
            pos += block.shape[0]

    def close(self) -> None:
        if not self.fh.closed:
            self.fh.close()

    def __enter__(self) -> "WavReader":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    def info(self) -> dict:
        return {
            "path": self.path,
            "samplerate": self.samplerate,
            "channels": self.channels,
            "encoding": self.encoding,
            "bits": self.bits,
            "frames": self.frames,
            "duration": self.duration,
            "data_bytes": self.data_size,
            "rf64": self.rf64,
            "format": self.fmt.describe(),
        }


def wav_info(path: str) -> dict:
    with WavReader(path) as r:
        return r.info()


def read_wav(path: str, *, max_seconds: float | None = None) -> tuple[np.ndarray, int]:
    """读取整个 WAV，返回 ``(data, samplerate)``；``max_seconds`` 可只读前段。"""
    with WavReader(path) as r:
        n = None if max_seconds is None else int(max_seconds * r.samplerate)
        return r.read(0, n), r.samplerate


def write_wav(path: str, data: np.ndarray, samplerate: int, *,
              bit_depth: int | str = 24, dither: bool = True,
              rf64: bool = False) -> dict:
    """一次性写出 WAV（离线后期处理用）。"""
    x = np.asarray(data, dtype=np.float32)
    if x.ndim == 1:
        x = x[:, None]
    encoding, bits = _fmt_for_bitdepth(bit_depth)
    fmt = WavFormat(int(samplerate), int(x.shape[1]), encoding, bits)
    w = WavWriter(path, fmt, dither=dither, rf64=rf64)
    try:
        # 分块写，避免超大数组一次性编码造成内存峰值
        step = 1 << 18
        for i in range(0, x.shape[0], step):
            w.write(x[i:i + step])
    finally:
        stats = w.close()
    stats["format"] = fmt.describe()
    return stats