"""Self-delimiting live transport: 46 payload bits in each 48-bit record.

Each six-byte frame has a leading one bit and a zero padding bit. A footer
starts with a zero high bit and holds a 31-bit sample count plus CRC32. The
48-byte model header and eight-byte footer add 56 fixed bytes. Payload is
2300 bit/s; framing adds 100 bit/s. CRC is verified at EOS, so earlier indices
are provisional. No seek, full-file pass or duration-sized buffer is needed.
"""
from dataclasses import dataclass
import struct
import zlib

import numpy as np

from tinyasr.codec_bitstream import BITS_PER_FRAME, DEFAULT_MAX_FRAMES, HOP, LEVELS, SAMPLE_RATE, WIDTHS, model_bytes

MAGIC=b'TCLIVE1\0'
HEADER=struct.Struct('<8sBBHI32s')
FOOTER_PREFIX=struct.Struct('>I')
CRC=struct.Struct('<I')
FOOTER_BYTES=FOOTER_PREFIX.size+CRC.size
RECORD_BYTES=6


def frame_limit(value):
    if isinstance(value,bool) or not isinstance(value,int) or not 1<=value<=DEFAULT_MAX_FRAMES:
        raise ValueError('Invalid live transport frame limit')
    return value


def empty_codes():return np.empty((4,6,0),dtype=np.int64)


@dataclass(frozen=True)
class LiveEnd:
    samples:int
    frames:int
    file_bytes:int
    model_sha256:str

    def accounting(self):
        seconds=self.samples/SAMPLE_RATE
        return dict(samples=self.samples,frames=self.frames,seconds=seconds,
                    payload_bits=self.frames*BITS_PER_FRAME,framing_bits=self.frames*2,
                    file_bytes=self.file_bytes,fixed_container_bytes=HEADER.size+FOOTER_BYTES,
                    file_bits_per_second=self.file_bytes*8/seconds,
                    model_sha256=self.model_sha256,format='tinycall-live-v1')


class LiveBitWriter:
    def __init__(self,model_sha256,max_frames=DEFAULT_MAX_FRAMES):
        self.max_frames=frame_limit(max_frames)
        self.header=HEADER.pack(MAGIC,1,1,BITS_PER_FRAME,SAMPLE_RATE,model_bytes(model_sha256))
        self.checksum=zlib.crc32(self.header);self.frames=0;self.finished=False

    @property
    def retained_bit_reservoir_bytes(self):return 0

    def push(self,codes):
        if self.finished:raise ValueError('Live writer already finished')
        codes=np.asarray(codes)
        if codes.ndim!=3 or codes.shape[:2]!=(4,6) or codes.dtype.kind not in 'iu':
            raise ValueError('Expected integer scalar codes [4,6,frames]')
        if codes.shape[-1]+self.frames>self.max_frames:raise ValueError('Live frame limit exceeded')
        if np.any(codes<0) or np.any(codes>=LEVELS[:,:,None]):raise ValueError('Scalar index outside its grid')
        output=bytearray()
        for index in range(codes.shape[-1]):
            word=0
            for value,width in zip(codes[:,:,index].reshape(-1),WIDTHS):word=(word<<int(width))|int(value)
            output.extend(((1<<47)|(word<<1)).to_bytes(RECORD_BYTES,'big'));self.frames+=1
        self.checksum=zlib.crc32(output,self.checksum)
        return bytes(output)

    def finish(self,samples):
        if self.finished:raise ValueError('Live writer already finished')
        if isinstance(samples,bool) or not isinstance(samples,(int,np.integer)) or not 0<samples<2**31:
            raise ValueError('Expected positive 31-bit sample count')
        if not self.frames or self.frames!=(int(samples)+HOP-1)//HOP:
            raise ValueError('Sample count does not match received frames')
        prefix=FOOTER_PREFIX.pack(int(samples));self.finished=True
        return prefix+CRC.pack(zlib.crc32(prefix,self.checksum))


class LiveBitReader:
    def __init__(self,expected_model_sha256,max_frames=DEFAULT_MAX_FRAMES):
        self.expected_model=model_bytes(expected_model_sha256);self.max_frames=frame_limit(max_frames)
        self.header_buffer=bytearray();self.record_buffer=bytearray();self.record_size=0
        self.header_complete=False;self.finished=False;self.failed=False;self.end=None
        self.frames=0;self.received_bytes=0;self.checksum=0

    @property
    def retained_transport_buffer_bytes(self):return len(self.header_buffer)+len(self.record_buffer)

    def _record(self,record):
        if self.record_size==RECORD_BYTES:
            value=int.from_bytes(record,'big')
            if value&1:raise ValueError('Nonzero live frame padding bit')
            if self.frames>=self.max_frames:raise ValueError('Live frame limit exceeded')
            word=(value>>1)&((1<<BITS_PER_FRAME)-1);shift=BITS_PER_FRAME;values=[]
            for width in WIDTHS:
                shift-=int(width);values.append((word>>shift)&((1<<int(width))-1))
            self.frames+=1;self.checksum=zlib.crc32(record,self.checksum)
            return np.array(values,dtype=np.int64).reshape(4,6)
        prefix=record[:FOOTER_PREFIX.size];samples=FOOTER_PREFIX.unpack(prefix)[0]
        if not 0<samples<2**31 or self.frames!=(samples+HOP-1)//HOP:
            raise ValueError('Live frame/sample counts disagree')
        if zlib.crc32(prefix,self.checksum)!=CRC.unpack(record[FOOTER_PREFIX.size:])[0]:
            raise ValueError('Live checksum mismatch')
        self.end=LiveEnd(samples,self.frames,self.received_bytes,self.expected_model.hex())
        return None

    def push(self,data):
        if self.finished or self.failed:raise ValueError('Live reader is closed or failed')
        if not isinstance(data,(bytes,bytearray,memoryview)):raise ValueError('Expected incoming bytes')
        try:data=memoryview(data).cast('B')
        except (TypeError,ValueError) as error:raise ValueError('Expected contiguous bytes') from error
        limit=HEADER.size+FOOTER_BYTES+self.max_frames*RECORD_BYTES
        if self.received_bytes+len(data)>limit:
            self.failed=True;raise ValueError('Live byte limit exceeded')
        self.received_bytes+=len(data)
        try:
            if self.end is not None:
                if len(data):raise ValueError('Unexpected trailing live data')
                return empty_codes()
            if not self.header_complete:
                need=HEADER.size-len(self.header_buffer)
                used=min(need,len(data));self.header_buffer.extend(data[:used]);data=data[used:]
                if len(self.header_buffer)<HEADER.size:return empty_codes()
                magic,version,flags,bits,rate,fingerprint=HEADER.unpack(self.header_buffer)
                if (magic,version,flags,bits,rate)!=(MAGIC,1,1,BITS_PER_FRAME,SAMPLE_RATE):raise ValueError('Unsupported live header')
                if fingerprint!=self.expected_model:raise ValueError('Live stream requires a different model checkpoint')
                self.checksum=zlib.crc32(self.header_buffer);self.header_buffer=bytearray();self.header_complete=True
            frames=[];position=0
            while position<len(data):
                if self.end is not None:raise ValueError('Unexpected trailing live data')
                if not self.record_buffer:self.record_size=RECORD_BYTES if data[position]&128 else FOOTER_BYTES
                used=min(self.record_size-len(self.record_buffer),len(data)-position)
                self.record_buffer.extend(data[position:position+used]);position+=used
                if len(self.record_buffer)==self.record_size:
                    value=self._record(bytes(self.record_buffer));self.record_buffer=bytearray()
                    if value is not None:frames.append(value)
            return np.stack(frames,-1) if frames else empty_codes()
        except (ValueError,struct.error):
            self.failed=True;raise

    def finish(self):
        if self.finished or self.failed:raise ValueError('Live reader is closed or failed')
        if not self.header_complete or self.record_buffer or self.end is None:
            self.failed=True;raise ValueError('Truncated live stream')
        self.finished=True;return self.end