#!/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
import cython, typing, gzip

# see warc.pxd for function signatures

INFO: int =  0
RESP: int = 1
REQ: int =  2
META: int = 3
REVISIT: int = 4

BUFSIZE: int = 16 * 1024 * 1024
BUFMIN: int = 3 * 1024 * 1024 # 1.5MiB, will need to be increased
                             # to 5.5MiB for CC-MAIN-2025-13 (Mar) and thereafter

HDRMAX: int = 0  # will grow
RECORDMAX: int = 0  # will grow

def warc(filename: str,
         callback: typing.Callable[[bytes, typing.ByteString, int], typing.BinaryIO],
         types: typing.List[int] = [RESP], whole: bool = False, parts: int = 7,
         debug: bool = False):
  '''parts is a bit-mask:
     1 for warc header;
     2 for req/resp HTTP header, warcinfo/metadata features;
     4 for req/resp body'''
  # Not currently trying to depend on this, but I believe that
  #   warcinfo: record-headers+1bl+crawl-headers+2bl
  #   request: record-headers+1bl+HTTP-headers+3bl
  #   response: record-headers+1bl+HTTP-headers+[1bl or 2bl]+HTTP-body+1bl
  #   metadata: record-headers+1bl+metadata-headers+3bl
  global BUFSIZE, HDRMAX, BUFMIN, RECORDMAX
  _out: typing.BinaryIO
  # should do some sanity checking wrt parts and types
  stream: typing.BinaryIO
  with gzip.open(filename, 'r') as fh:
    try:
      fh.read(1)
    except gzip.BadGzipFile:
      stream = open(filename,'rb',0)
    else:
      stream = igzip.IGzipFile(filename=filename)
  buf: char[::1] = bytearray(BUFSIZE)
  bufView: char[::1] = memoryview(buf)
  fpos: int = 0
  bp: int = 0
  bl: int = stream.readinto(buf)
  n: int = 0
  done: bool = bl < BUFSIZE 
  while buf.startswith(b'\r\n',bp):
    bp+=2
  while not (done and bp >= bl):
    start_1: int = bp
    if not buf.startswith(b'WARC/1.0\r\n',bp):
      breakpoint()
      raise ValueError("Not a WARC file? In %s at %s of %s (%s): %s[%s]"%(filename,
                                                                   bp,bl,fpos,
         (buf[bp:min(bl,bp+20)] if bp<bl else buf[bl-20:bl]).decode('latin-1'),
                                                                     bl-bp))
    bp += 10
    n += 1
    wtype: int = -1
    length: int = 0
    tr: bytes = b'' # 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 at the end
      eol = buf.index(b'\r\n', bp)
      if buf.startswith(b"Content-Length: ",bp):
        length=wl=int(bufView[bp+16:eol])
      if buf.startswith(b"WARC-Truncated: ",bp):
        if bp+16==eol-2:
          tr = b"EMPTY"
        else:
          tr = bytes(bufView[bp+16:eol-2])
      elif buf.startswith(b'WARC-Type: ',bp):
        if buf.startswith(b's',bp+13):
          wtype = RESP
        elif buf.startswith(b'q',bp+13):
          wtype = REQ
        elif buf.startswith(b'm',bp+11):
          wtype = META
        elif buf.startswith(b'w',bp+11):
          wtype = INFO
        elif buf.startswith(b'v',bp+13):
          wtype = REVISIT
        else:
          raise ValueError("Unknown WARC-Type: %s in %s at %s"%(
                             bytes(bufView[bp+11:eol-2]),filename,
                             fpos-(bl-bp)))
      bp=eol+2
    # record header done
    if (hl:=(bp - start_1)) > HDRMAX:
      HDRMAX = hl
    #if done:
    #  if (bp+length)>bl:
    #    raise ValueError("Done but need more! %s + %s > %s in %s"%(bp,
    #                     length,bl,filename))
    if (wtype in types):
      # Output whole or part 1 as required
      if whole:
        _out=callback(wtype,bufView[start_1:bp+length],7)
      else:
        if (parts & 1):
          bp = eol+2
          _out=callback(wtype,bufView[start_1:bp],1)
        if parts!=1:
          while buf.startswith(b'\r\n',bp):
            bp+=2
          start_2=bp
          eob=bp+length
          while buf.startswith(b'\r\n',eob-2):
            eob-=2
          # Only output parts (2 = HTTP header, 4 = body) that are wanted
          if parts & 2:
            if wtype == RESP or wtype == REQ :
              # request and response have http headers
              eo2=buf.index(b'\r\n\r\n',start_2)
              _out=callback(wtype,bufView[start_2:eo2+2],2)
            else:
              # rest of the part
              _out=callback(wtype,bufView[start_2:eob],2)
          if parts & 4:
            raise ValueError("Not implemented: body part (4): %s"%parts)
    #bp += length
    #if buf[bp] != 13:
    #  # Why does this sometimes happen, e.g. when doing
    #  python3 ~/lib/python/cc/test_warc.py 4 /beegfs/common_crawl/CC-MAIN-2019-35/1566027313501.0/orig/crawldiagnostics/CC-MAIN-20190817222907-20190818004907-00000.warc.gz
    #  at a point where bp+length is 11018, looking at >\n\r\n
    #  bp += 1 [doesn't work]
    bp = buf.index(b'\r\n',bp+length)
    # check if refill needed
    rl: int
    if (rl := (bp - start_1)) > RECORDMAX:
      RECORDMAX = rl
    keepLen: int
    if (not done) and (keepLen := bl - bp) < BUFMIN:
      # we need to shift and read more
      buf[0:keepLen]=bufView[bp:bl]
      with memoryview(buf)[keepLen:BUFSIZE] as xBuf:
        nb=stream.readinto(xBuf)
      bl = keepLen+nb
      done = bl < BUFSIZE 
      bp = 0
    while buf.startswith(b'\r\n',bp):
      bp+=2
    #print('end of loop',wtype,start_1,bp,eol,length,bl,file=sys.stderr)
  print('%d records, max record: %d, max header: %d'%(n, RECORDMAX, HDRMAX),
        file=sys.stderr)

import zlib, gzip, struct

def decompOneBlock(fp: io.BytesIO):
    """Decompress one block of a gzip compressed stream in one shot.
    Return the decompressed string and the stream repositioned
      at the start of the next block.
    """
    data = fp.getbuffer()
    if gzip._read_gzip_header(fp) is None:
        return b""
    # Use a zlib raw deflate compressor
    do = zlib.decompressobj(wbits=-zlib.MAX_WBITS)
    # Read all the data except the header
    decompressed = do.decompress(data[fp.tell():])
    if not do.eof or len(do.unused_data) < 8:
        raise EOFError("Compressed file ended before the end-of-stream "
                       "marker was reached")
    crc, length = struct.unpack("<II", do.unused_data[:8])
    if crc != zlib.crc32(decompressed):
        raise BadGzipFile("CRC check failed")
    if length != (len(decompressed) & 0xffffffff):
        raise BadGzipFile("Incorrect length of data produced")
    eob = 8
    unused = len(do.unused_data) - 8
    while unused > 0 and do.unused_data[eob] == 0:
        eob += 1
        unused -= 1
    if unused > 0:
        fp.seek(-unused, 2)
    return decompressed