#!/usr/bin/env python3
"""Minimal jyoti usage: encode a 16 kHz mono file to a 2.3 kbps stream and decode it back, in process.

    python example.py input_16k_mono.wav output.wav [--checkpoint codec.pt]

The stream is bound to the checkpoint's SHA-256; decoding with a different checkpoint raises an error.
"""
import argparse
import hashlib
from pathlib import Path
import sys

import soundfile as sf
import torch

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
from tinyasr.codec_factory_v3 import load_codec_checkpoint  # noqa: E402
from tinyasr.codec_live_v1 import LiveEncoder  # noqa: E402
from tinyasr.codec_phase_streaming_v1 import PhaseLiveDecoder  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('input', type=Path)
    ap.add_argument('output', type=Path)
    ap.add_argument('--checkpoint', type=Path, default=ROOT / 'codec.pt')
    args = ap.parse_args()
    if args.output.exists():
        raise FileExistsError(args.output)
    audio, rate = sf.read(args.input, dtype='float32', always_2d=False)
    if rate != 16000 or audio.ndim != 1:
        raise ValueError(f'Expected 16 kHz mono input, got {rate} Hz with shape {audio.shape}')
    torch.set_num_threads(1)
    model_sha256 = hashlib.sha256(args.checkpoint.read_bytes()).hexdigest()
    codec = load_codec_checkpoint(args.checkpoint)

    # Encode in 20 ms (320-sample) blocks, as a live sender would.
    encoder = LiveEncoder(codec, model_sha256)
    stream = bytearray(encoder.start())
    x = torch.from_numpy(audio)
    for start in range(0, x.numel(), 320):
        stream += encoder.push(x[start:start + 320])
    stream += encoder.finish()

    # Decode the byte stream; finish() checks the footer and CRC.
    decoder = PhaseLiveDecoder(codec, model_sha256)
    parts = [decoder.push(bytes(stream))]
    tail, end = decoder.finish()
    parts.append(tail)
    decoded = torch.cat(parts).numpy()
    sf.write(args.output, decoded, 16000, subtype='FLOAT')
    seconds = audio.size / 16000
    print(f'{seconds:.2f} s audio -> {len(stream)} bytes ({8 * len(stream) / seconds:.0f} bit/s incl. framing); '
          f'decoded {decoded.size} samples -> {args.output}')


if __name__ == '__main__':
    main()