#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Network-free streaming receiver for one Phase-1.0DM result."""
from __future__ import annotations
from dataclasses import dataclass
import hashlib,os
from pathlib import Path
from phase10dm_snapshot_protocol import BEGIN,END,FRAME,MAGIC,MAX_BYTES,ProtocolError,_frame

@dataclass(frozen=True)
class ReceivedSnapshot:
 path:Path;size:int;sha256:str;device:int;inode:int;mtime:int

class SnapshotReceiver:
 def __init__(self,path:Path):
  if not isinstance(path,Path) or path.exists():raise ProtocolError("exclusive output required")
  self.path=path;self.stream=path.open("xb");self.buffer=bytearray();self.state="HEADER";self.expected=None;self.written=0;self.digest=hashlib.sha256();self.begin=None
 def feed(self,chunk:bytes):
  if self.state not in {"HEADER","PAYLOAD","TERMINAL"} or not isinstance(chunk,bytes):raise ProtocolError("receiver state")
  self.buffer.extend(chunk)
  if self.state=="HEADER" and len(self.buffer)>=64:
   self.begin=_frame(bytes(self.buffer[:64]));del self.buffer[:64];kind,status,saved,size,sent,*_=self.begin[2:]
   if kind!=BEGIN or status or saved or sent or not 0<size<=MAX_BYTES:self.abort();raise ProtocolError("begin invalid")
   self.expected=size;self.state="PAYLOAD"
  if self.state=="PAYLOAD":
   take=min(len(self.buffer),self.expected-self.written)
   if take:
    data=bytes(self.buffer[:take]);del self.buffer[:take];self.stream.write(data);self.digest.update(data);self.written+=take
   if self.written==self.expected:self.state="TERMINAL"
  if self.state=="TERMINAL" and len(self.buffer)>64:self.abort();raise ProtocolError("trailing bytes")
 def finish(self)->ReceivedSnapshot:
  if self.state!="TERMINAL" or len(self.buffer)!=64 or self.begin is None:self.abort();raise ProtocolError("stream incomplete")
  end=_frame(bytes(self.buffer));kind,status,saved,size,sent,dev,ino,mtime=end[2:];_,_,_,bsize,_,bdev,bino,bmtime=self.begin[2:]
  if kind!=END or status or saved or size!=bsize or sent!=bsize or (dev,ino,mtime)!=(bdev,bino,bmtime):self.abort();raise ProtocolError("terminal invalid")
  self.stream.flush();os.fsync(self.stream.fileno());self.stream.close();self.state="SEALED"
  with self.path.open("rb") as check:
   if check.read(16)!=b"SQLite format 3\x00":raise ProtocolError("not sqlite")
  return ReceivedSnapshot(self.path,self.written,self.digest.hexdigest(),dev,ino,mtime)
 def abort(self):
  if not self.stream.closed:self.stream.flush();self.stream.close()
  self.state="ABORTED"