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

if (debug:=(sys.argv[1]=='-d')):
  sys.argv.pop(1)

def warc(callback,types=['response'],parts=7):
  types=[(t if isinstance(t,bytes) else bytes(t,'utf8')) for t in types]
  nb=0
  stream=open(sys.argv[1],'rb',0)
  bufsize=128*1024*1024
  buf=bytearray(128*1024*1024)
  l=b'\r\n'
  while True:
    bp=0
    while l==b'\r\n':
      l=stream.readline()
      nb+=(ln:=len(l))
    if l!=b'WARC/1.0\r\n':
      raise ValueError("Not a WARC file? At %s: %s[%s]"%(nb-len(l),
                                                         l.decode('latin-1'),len(l)))
    wtype=None
    length=None
    state=1
    tr=None # Was this record truncated?
    while l!=b'\r\n':
      if parts & 1:
        buf[bp:(bp:=bp+ln)]=l
      l=stream.readline()
      nb+=(ln:=len(l))
      # WARC header
      if l.startswith(b"Content-Length: "):
        length=wl=int(l[16:].rstrip())
      elif l.startswith(b"WARC-Truncated: "):
        tr=l[16:].rstrip()
        tr="EMPTY" if tr=="" else tr
      elif l.startswith(b'WARC-Type: '):
        wtype = l[11:-2]
    start_2=bp
    if (wtype in types) and (parts & 1):
      if parts!=1:
        buf[bp:(bp:=bp+ln)]=l
        start_2=bp
      if parts!=7:
        callback(wtype,buf[:start_2],1)
    else:
      start_2=0
    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
    bv=memoryview(buf)[start_2:start_2+length]
    if wtype in types:
      if parts==7:
        callback(wtype,memoryview(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=0
      state=2
      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,bv[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=rec_text.tell()
              #print(balance,bl,wl,ll,ll-balance,file=sys.stderr)
              # Output whatever is left
              if parts & 4:
                callback(wtype,bv[balance:balance+wl],4)
              state=1
              
          L_start=rec_text.tell()
OUT=open(sys.stdout.fileno(),'wb')

import re
LMPAT=re.compile(b'^Last-Modified: (.*?)\r',re.MULTILINE)

def showmeLMH(wtype,buf,part=2):
  m=LMPAT.search(buf.tobytes(order='A'))
  if m:
    OUT.write(m[1])
  OUT.write(b'\n')

def showme(wtype,buf,part):
  if debug:
    breakpoint()
  OUT.write(b"%d\n%b"%(part,buf))

#warc(showmeLMH,[b'response'],2)

#warc(showme,[b'response','warcinfo','request','metadata'],int(sys.argv[2]))

warc(showme,[b'response'],int(sys.argv[2]))