#!/usr/bin/env python3
"""Unknown-duration waveform/byte streaming through files or raw stdin/stdout."""
import argparse
from contextlib import ExitStack
import hashlib
import json
import os
from pathlib import Path
import sys
import time

import numpy as np
import soundfile as sf
import torch

ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT))
from tinyasr.codec_factory_v3 import load_codec_checkpoint
from tinyasr.codec_live_v1 import LiveDecoder,LiveEncoder
from tinyasr.codec_phase_streaming_v1 import PhaseLiveDecoder
from tinyasr.codec_phase_residual_v1 import PhaseResidualDecoder


def chunks(source,size):
    read=getattr(source,'read1',source.read)
    while True:
        data=read(size)
        if not data:break
        yield data


def raw_audio(source,format,block_samples):
    width=2 if format=='s16le' else 4
    pending=b''
    for data in chunks(source,block_samples*width):
        data=pending+data;complete=len(data)//width*width;pending=data[complete:]
        if complete:
            x=np.frombuffer(data[:complete],dtype='<i2' if width==2 else '<f4').astype('float32')
            if width==2:x/=32768
            yield torch.from_numpy(x)
    if pending:raise ValueError('Raw input ended inside a PCM sample')


def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('operation',choices=['encode','decode'])
    ap.add_argument('--checkpoint',type=Path,required=True)
    ap.add_argument('--input',required=True)
    ap.add_argument('--output',required=True)
    ap.add_argument('--input-format',choices=['audio','s16le','f32le'],default='audio')
    ap.add_argument('--output-format',choices=['wav','f32le'],default='wav')
    ap.add_argument('--block-samples',type=int,default=320)
    ap.add_argument('--threads',type=int,default=1)
    args=ap.parse_args()
    if not 1 <= args.threads <= 4:raise ValueError('Threads must be between one and four')
    if args.block_samples<1:raise ValueError('Block size must be positive')
    if args.input=='-' and args.operation=='encode' and args.input_format=='audio':
        raise ValueError('Raw live stdin requires --input-format s16le or f32le')
    if args.output=='-' and args.operation=='decode' and args.output_format!='f32le':
        raise ValueError('Live stdout requires --output-format f32le')
    target=Path(args.output) if args.output!='-' else None
    partial=target.with_name(target.name+'.partial') if target else None
    if target and (target.exists() or partial.exists()):raise FileExistsError('Output or its partial file exists')
    torch.set_num_threads(args.threads)
    fingerprint=hashlib.sha256(args.checkpoint.read_bytes()).hexdigest()
    codec=load_codec_checkpoint(args.checkpoint)
    started=time.monotonic()
    with ExitStack() as stack:
        if args.operation=='encode':
            sender=LiveEncoder(codec,fingerprint)
            output=stack.enter_context(partial.open('xb')) if target else sys.stdout.buffer
            output.write(sender.start());output.flush()
            if args.input_format=='audio':
                source=stack.enter_context(sf.SoundFile(args.input))
                if source.samplerate!=16000 or source.channels!=1:raise ValueError('Expected 16 kHz mono input')
                audio=(torch.from_numpy(x) for x in source.blocks(blocksize=args.block_samples,dtype='float32'))
            else:
                source=stack.enter_context(Path(args.input).open('rb')) if args.input!='-' else sys.stdin.buffer
                audio=raw_audio(source,args.input_format,args.block_samples)
            for x in audio:
                output.write(sender.push(x));output.flush()
            output.write(sender.finish());output.flush()
            samples=sender.audio.samples;state_bytes=sender.float_state_bytes
            result=dict(samples=samples,frames=sender.audio.frames,payload_bits=sender.audio.frames*46,
                        framing_bits=sender.audio.frames*2,wire_bytes=48+6*sender.audio.frames+8)
        else:
            receiver=(PhaseLiveDecoder if isinstance(codec.decoder,PhaseResidualDecoder) else LiveDecoder)(codec,fingerprint)
            source=stack.enter_context(Path(args.input).open('rb')) if args.input!='-' else sys.stdin.buffer
            if args.output_format=='wav':
                output=stack.enter_context(sf.SoundFile(partial,mode='w',samplerate=16000,channels=1,format='WAV',subtype='FLOAT'))
                write=lambda x:output.write(x.numpy())
            else:
                output=stack.enter_context(partial.open('xb')) if target else sys.stdout.buffer
                def write(x):
                    output.write(x.numpy().astype('<f4',copy=False).tobytes());output.flush()
            for data in chunks(source,4096):
                write(receiver.push(data))
                if args.input=='-' and receiver.transport.end is not None:break
            tail,end=receiver.finish();write(tail)
            samples=end.samples;state_bytes=receiver.float_state_bytes
            result=end.accounting()
    # File outputs become visible at the requested path only after complete
    # finalization/CRC checks. Raw stdout has already delivered provisional data.
    if target:
        os.link(partial,target)
        partial.unlink()
    elapsed=time.monotonic()-started
    result.update(operation=args.operation,elapsed_seconds=elapsed,rtf=elapsed/(samples/16000),
                  persistent_float_state_bytes=state_bytes,model_sha256=fingerprint,quality_qualified=False)
    print(json.dumps(result),file=sys.stderr)


if __name__=='__main__':main()