#!/usr/bin/env python3
'''Stream a warc format file, unzipping if necessary, invoking a
callback on each record.  Callback can be limited by WARC-Type, record
part'''

import sys,io
from isal import igzip

def warc(filename,callback,types=['response'],whole=False,parts=7,debug=False):
  types=[(t if isinstance(t,bytes) else bytes(t,'utf8')) for t in types]
  nb=0
  if filename.endswith(".gz"):
    stream=igzip.IGzipFile(filename=filename)
  else:
    stream=open(filename,'rb',0)
  bufSize=2*1024*1024
  hdrMax=16*1024
  buf=bytearray(bufSize)
  #with memoryview(buf)[:hdrMax] as hdrBuf:
  fpos=bl=stream.readinto(buf)
  bp=0
  done=False
  while True:
    while buf.startswith(b'\r\n',bp): # will Fail if buffer (nearly) empty 
      bp+=2
    start_1=bp
    if not buf.startswith(b'WARC/1.0\r\n',bp):
      if done and bl-bp==0:
        # really done
        return
      raise ValueError("Not a WARC file? At %s: %s[%s]"%(bp,
                       buf[bp:min(bl,bp+20)].decode('latin-1'), bl-bp))
    bp+=10
    wtype=None
    length=None
    state=1
    tr=None # Was this record truncated?
    while not buf.startswith(b'\r\n',bp):
      # there should always be enough in the buffer to complete this loop,
      #  because of the buffer update logic below
      eol=buf.index(b'\r\n',bp)+2
      if buf.startswith(b"Content-Length: ",bp):
        length=wl=int(buf[bp+16:eol-2])
      elif buf.startswith(b"WARC-Truncated: ",bp):
        tr=l[bp+16:eol-2]
        tr="EMPTY" if tr=="" else tr
      elif buf.startswith(b'WARC-Type: ',bp):
        wtype = bytes(buf[bp+11:eol-2])
      bp=eol
    bp=eol+2
    if (bp+length)>bl:
      if done:
        raise ValueError("Done but need more! %s + %s > %s",bp,length,bl)
      # Need more data
      if wtype in types:
        # we need to keep from start_1 to bl
        keepFrom=start_1
        keepLen=bl-keepFrom
        buf[0:keepLen]=buf[keepFrom,bl]
      else:
        # we can skip the rest of this part
        keepLen=0
        fpos=stream.seek(fpos+(pb+length-bl))
      spaceToFill=bufMax-keepLen
      with memoryview(buf)[keepLen:bufMax] as xBuf:
        nb=stream.readinto(xBuf)
      fpos+=nb
      bp=keepLen
      bl=keepLen+nb
      if nb<spaceToFill:
        done=True
      if wtype not in types:
        continue
    if (wtype in types):
      if whole:
        pass # buf[bp:(bp:=bp+ln)]=l @fixme
      elif (parts & 1):
        OUT=callback(wtype,buf[start_1:eol],1)
      if parts!=1:
        # everything from bv= goes here
        pass
      else:
        bp+=length
    print('end of loop',wtype,start_1,bp,eol,length,file=sys.stderr)
    #while not buf.startswith(b'\r\n',bp):
    OUT.write(b"=====\n")
    OUT.write(buf[0:100])
    if not buf[99]==10:
      OUT.write(b"\n")
    continue
    return
    bv=memoryview(buf)[start_2:start_2+length]
    ii=0
    while True and not stream.closed:
      if (i:=stream.readinto(bv))==0:
        break
      ii+=i
      if ii>=length:
        break
      bv=memoryview(buf)[start_2+ii:start_2+length]
    if ii!=length:
      raise ValueError("Chunk read losing, from %s got %s expected %s"%(nb,ii,length))
    nb+=length
    if wtype in types:
      if whole:
        callback(wtype,buf[0:start_2+length],7)
        continue
      # Only output parts (1 = WARC header, 2 = HTTP header, 4 = body) that are wanted
      bl=None # for HTTP Content-Length for the length of the body?
      L_start=start_2
      state=2
      bv=memoryview(buf)[start_2:start_2+length]
      with io.BytesIO(bv) as rec_text:
        for L in rec_text:
          if state==2:
            # HTTP header
            wl -= len(L)
            if not (L==b"" or L.startswith(b"\r")):
              # Non-empty, it's (a continuation of) a header
              if bl is None and L.startswith(b"Content-Length: "):
                bl=int(L[16:].rstrip())
            else:
              # Blank line, HTTP header is finished
              if parts & 2:
                callback(wtype,buf[start_2:start_2+L_start],2)
              state=4
              # The above is just for sanity, because we do _not_
              #  continue with the outer loop,
              #  since we can now block-output the entire rest of the
              #  input buffer.
              if bl is not None:
                if bl!=wl:
                  print("length mismatch: %s %s %s here: %s given: %s trunc: %s"%\
                        (length,offset,filename,wl,bl,tr),file=sys.stderr)
              # HTTP body
              balance=start_2+rec_text.tell()
              #print(balance,bl,wl,ll,ll-balance,file=sys.stderr)
              # Output whatever is left
              if parts & 4:
                callback(wtype,buf[balance:balance+wl],4)
              state=1
              
          L_start=rec_text.tell()