from __future__ import annotations

import struct
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
from encodec import EncodecModel

# =========================================================
# CONFIG
# =========================================================

INPUT_RAW = "music_1_5.raw"
OUTPUT_WAV = "reconstructed.wav"

BIT_DEPTH = 10
MAGIC = b"OCS1"

BW_TO_NQ = {
    1500: 2,
    3000: 4,
    6000: 8,
    12000: 16,
    24000: 32,
}

# =========================================================
# BIT UNPACKING
# =========================================================

def unpack_bits(data: bytes, bits_per_value: int = 10, expected_values: int | None = None) -> np.ndarray:
    """Unpack a little-endian bitstream into integers."""
    values = []
    bitstream = 0
    bitcount = 0
    mask = (1 << bits_per_value) - 1

    for b in data:
        bitstream |= b << bitcount
        bitcount += 8

        while bitcount >= bits_per_value:
            values.append(bitstream & mask)
            bitstream >>= bits_per_value
            bitcount -= bits_per_value

            if expected_values is not None and len(values) >= expected_values:
                return np.array(values, dtype=np.uint16)

    return np.array(values, dtype=np.uint16)

# =========================================================
# HEADER PARSING
# =========================================================

def parse_header(blob: bytes):
    """
    Returns:
        version, bandwidth_milli, sample_rate, n_q, num_frames, total_tokens, frame_lengths, payload_offset
    """
    offset = 0

    magic = blob[offset:offset + 4]
    offset += 4
    if magic != MAGIC:
        raise ValueError("Not a valid cassette file (bad magic).")

    version = struct.unpack_from("<H", blob, offset)[0]
    offset += 2

    bandwidth_milli = struct.unpack_from("<H", blob, offset)[0]
    offset += 2

    sample_rate = struct.unpack_from("<I", blob, offset)[0]
    offset += 4

    n_q = struct.unpack_from("<H", blob, offset)[0]
    offset += 2

    num_frames = struct.unpack_from("<I", blob, offset)[0]
    offset += 4

    total_tokens = struct.unpack_from("<I", blob, offset)[0]
    offset += 4

    frame_lengths = []
    for _ in range(num_frames):
        flen = struct.unpack_from("<I", blob, offset)[0]
        offset += 4
        frame_lengths.append(int(flen))

    return (
        version,
        bandwidth_milli,
        sample_rate,
        n_q,
        num_frames,
        total_tokens,
        frame_lengths,
        offset,
    )

# =========================================================
# MAIN
# =========================================================

def main() -> None:
    input_path = Path(INPUT_RAW)
    if not input_path.exists():
        raise FileNotFoundError(f"Missing input raw file: {input_path}")

    blob = input_path.read_bytes()

    (
        version,
        bandwidth_milli,
        sample_rate,
        n_q,
        num_frames,
        total_tokens,
        frame_lengths,
        payload_offset,
    ) = parse_header(blob)

    if bandwidth_milli not in BW_TO_NQ:
        raise ValueError(f"Unsupported bandwidth in file: {bandwidth_milli / 1000:.3f} kbps")

    expected_n_q = BW_TO_NQ[bandwidth_milli]
    if n_q != expected_n_q:
        raise ValueError(f"Header n_q mismatch: file has {n_q}, expected {expected_n_q}")

    payload = blob[payload_offset:]

    tokens = unpack_bits(payload, bits_per_value=BIT_DEPTH, expected_values=total_tokens)

    if len(tokens) < total_tokens:
        raise ValueError(f"Token stream truncated: got {len(tokens)}, expected {total_tokens}")

    tokens = tokens[:total_tokens]

    print(f"Version: {version}")
    print(f"Bandwidth: {bandwidth_milli / 1000:.3f} kbps")
    print(f"Sample rate: {sample_rate}")
    print(f"Codebooks: {n_q}")
    print(f"Frames: {num_frames}")
    print(f"Tokens: {len(tokens)}")

    encoded_frames = []
    idx = 0

    for flen in frame_lengths:
        count = n_q * flen
        frame_tokens = tokens[idx:idx + count]
        idx += count

        if len(frame_tokens) != count:
            raise ValueError("Token stream ended early while rebuilding frames.")

        codes = torch.tensor(frame_tokens, dtype=torch.long).reshape(1, n_q, flen)
        encoded_frames.append((codes, None))

    print("Loading EnCodec model...")
    model = EncodecModel.encodec_model_24khz()
    model.set_target_bandwidth(bandwidth_milli / 1000.0)
    model.eval()

    print("Decoding audio...")
    with torch.inference_mode():
        decoded = model.decode(encoded_frames)

    decoded = decoded.squeeze(0).cpu().numpy().T  # [T, C]

    sf.write(OUTPUT_WAV, decoded, sample_rate)
    print(f"Saved: {OUTPUT_WAV}")

if __name__ == "__main__":
    main()