from __future__ import annotations

import os
import struct
import subprocess
import tempfile
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
from encodec import EncodecModel
from encodec.utils import convert_audio

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

INPUT_AUDIO = "input_music.mp3"      # can be mp3, wav, etc.
OUTPUT_RAW = "music.raw"

BANDWIDTH = 1.5    # 1.5, 3, 6, 12, 24
BIT_DEPTH = 10
MODEL_SR = 24000

BW_TO_NQ = {
    1.5: 2,
    3.0: 4,
    6.0: 8,
    12.0: 16,
    24.0: 32,
}

MAGIC = b"OCS1"
VERSION = 1

# =========================================================
# BIT PACKING
# =========================================================

def pack_bits(values: np.ndarray, bits_per_value: int = 10) -> bytes:
    """Pack uint integers into a compact little-endian bitstream."""
    bitstream = 0
    bitcount = 0
    out = bytearray()

    mask = (1 << bits_per_value) - 1

    for v in values:
        bitstream |= (int(v) & mask) << bitcount
        bitcount += bits_per_value

        while bitcount >= 8:
            out.append(bitstream & 0xFF)
            bitstream >>= 8
            bitcount -= 8

    if bitcount > 0:
        out.append(bitstream & 0xFF)

    return bytes(out)

# =========================================================
# HELPERS
# =========================================================

def decode_input_to_temp_wav(input_path: str, temp_wav_path: str) -> None:
    """
    Decode any ffmpeg-supported input to mono 24 kHz WAV.
    This avoids torchaudio/TorchCodec issues and keeps input handling simple.
    """
    cmd = [
        "ffmpeg",
        "-y",
        "-i", input_path,
        "-ac", "1",
        "-ar", str(MODEL_SR),
        temp_wav_path,
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

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

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

    if BANDWIDTH not in BW_TO_NQ:
        raise ValueError(f"Unsupported bandwidth: {BANDWIDTH}")

    n_q = BW_TO_NQ[BANDWIDTH]

    print("Loading EnCodec model...")
    model = EncodecModel.encodec_model_24khz()
    model.set_target_bandwidth(BANDWIDTH)
    model.eval()

    with tempfile.TemporaryDirectory() as td:
        temp_wav = str(Path(td) / "input_mono_24k.wav")

        print("Decoding input audio to temporary WAV...")
        decode_input_to_temp_wav(str(input_path), temp_wav)

        print("Loading temporary WAV...")
        wav_np, sr = sf.read(temp_wav, always_2d=True)

        # soundfile gives [samples, channels]; EnCodec expects [channels, samples]
        wav_np = wav_np.T
        wav = torch.from_numpy(wav_np).float()

        wav = convert_audio(
            wav,
            sr,
            model.sample_rate,
            model.channels
        )

        wav = wav.unsqueeze(0)  # [B, C, T]

        print("Encoding audio...")
        with torch.inference_mode():
            encoded_frames = model.encode(wav)

        frame_lengths = []
        all_tokens = []

        for codes, scale in encoded_frames:
            # codes: [B, n_q, T]
            if codes.shape[1] != n_q:
                raise RuntimeError(
                    f"Unexpected number of codebooks. Got {codes.shape[1]}, expected {n_q}."
                )

            frame_lengths.append(int(codes.shape[-1]))

            flat = codes.squeeze(0).cpu().numpy().astype(np.uint16).reshape(-1)
            all_tokens.append(flat)

        tokens = np.concatenate(all_tokens) if all_tokens else np.array([], dtype=np.uint16)

        print(f"Frame count: {len(frame_lengths)}")
        print(f"Total tokens: {len(tokens)}")

        packed = pack_bits(tokens, bits_per_value=BIT_DEPTH)

        # =====================================================
        # FILE FORMAT
        # =====================================================
        # magic[4]              = b'OCS1'
        # version uint16        = 1
        # bandwidth_milli uint16= 3000 for 3.0 kbps
        # sample_rate uint32    = 24000
        # n_q uint16            = 4 at 3 kbps
        # num_frames uint32
        # total_tokens uint32
        # frame_lengths[num_frames] as uint32 each
        # packed token bitstream
        # =====================================================

        header = bytearray()
        header += MAGIC
        header += struct.pack("<H", VERSION)
        header += struct.pack("<H", int(BANDWIDTH * 1000))
        header += struct.pack("<I", model.sample_rate)
        header += struct.pack("<H", n_q)
        header += struct.pack("<I", len(frame_lengths))
        header += struct.pack("<I", int(len(tokens)))

        for flen in frame_lengths:
            header += struct.pack("<I", int(flen))

        with open(OUTPUT_RAW, "wb") as f:
            f.write(header)
            f.write(packed)

    size_kb = os.path.getsize(OUTPUT_RAW) / 1024
    print(f"Saved: {OUTPUT_RAW}")
    print(f"Final size: {size_kb:.2f} KB")

if __name__ == "__main__":
    main()