#!/usr/bin/env python3
'''Extract request records from Common Crawl WARC-format files
given length, offset and filename triples.
Input one triple on command line, or
triples from stdin as tab-delimited lines
or complete cdx index lines.
In all cases by 'filename' is meant crawlid/segmentid/filename

Note that if no output flag(s) is/are given, the whole WARC record will be output, more efficiently than would be the case if -whb is given.'''

import sys, argparse, regex, os, shutil, io, gzip, time
from isal import igzip
#from subprocess import Popen, PIPE
#import asyncio

HACK_USAGE=regex.compile('\[-x\]\n\s*\[length\] \[offset\] \[filename\]')
BINOUT=sys.stdout.buffer
FPAT="/%s/%s/orig/warc/%s"

class HackFormat(argparse.RawDescriptionHelpFormatter):
  def format_help(self):
    global FOO
    FOO=argparse.RawDescriptionHelpFormatter.format_help(self)
    return HACK_USAGE.sub('\n             [ ( -x | length offset filename ) ]',
                          FOO)

def process(options,buf,root,filename,offset,length,whole):
  rfn=root+filename
  if root!="/beegfs/common_crawl":
    if not os.path.exists(rfn):
      if not os.path.exists(os.path.dirname(rfn)):
        os.makedirs(os.path.dirname(rfn))
      with io.FileIO('/beegfs/common_crawl'+filename,'r') as infile, \
              io.FileIO(rfn,'w') as outfile:
        #shutil.copyfileobj(infile,outfile,128*1024*1024)
        while True:
          l=infile.readinto(buf)
          if l==0:
            break
          outfile.write(memoryview(buf)[:l])
  file=open(rfn,'rb',0)
  if whole:
    # try external unzip using Popen
    file.seek(offset)
    bv=memoryview(buf)[:length]
    nb=file.readinto(bv)
    file.close()
    if nb!=length:
      print("losing",file.name,length,nb,file=sys.stderr)
    if options.zipped:
      BINOUT.write(bv)
    else:
      gzip_chunk = io.BytesIO(bv)
      uv=memoryview(buf)[length:]
      #clear_bytes=io.BytesIO(uv)
      with igzip.IGzipFile(fileobj=gzip_chunk) as gzip_fin:
        while True:
          l=gzip_fin.readinto(uv)
          if not l:
            break
          BINOUT.write(memoryview(uv)[:l])

def main():
  parser = argparse.ArgumentParser(
    description='''Extract records from warc files given length, offset and file triples.
  Input one triple on command line, or
  triples from stdin as tab-delimited lines
  or complete cdx index lines.
  In all cases by 'filename' is meant crawlid/segmentid/filename''',
    epilog='''Note that if no output flag(s) is/are given,
  the whole WARC record will be output, more efficiently than
  would be the case if all three flags were given.''',
    add_help=False,
    conflict_handler='resolve',
    formatter_class=HackFormat
    )

  parser.add_argument('--help',help='Show help',action='help')
  parser.add_argument('-d','--debug',help='Debug output',action='store_true')
  parser.add_argument('-w','--warc',help='output WARC headers',
                      action='store_true')
  parser.add_argument('-h','--headers',help='output HTTP headers',
                      action='store_true')
  parser.add_argument('-b','--body',help='output HTTP body',
                      action='store_true')
  parser.add_argument('-c','--cmd',help='pipes each result thru CMD')
  parser.add_argument('-r','--root',nargs='?',
                  help='File path root, create a copy there if necessary',
                  default='/beegfs/common_crawl'),
  parser.add_argument('-z','--zipped',
                      help="output raw gzipped record, ignored if any of -bhw supplied",
                      action='store_true')
  sg=parser.add_mutually_exclusive_group()
  sg.add_argument('-x','--index',
                      help='take lines of triples from a cdx index file as input',
                      action='store_true')
  sg.add_argument('length',type=int,
                  help='length in bytes of gzipped record',
                  nargs='?')
  parser.add_argument('offset',type=int,
                      help='start position in bytes of gzipped record',
                      nargs='?')
  parser.add_argument('filename',
                      help='pathname of gzipped Common Crawl WARC-format file',
                      nargs='?')
  # Hack the order of optional and positional in the help output
  parser._action_groups.sort(key=lambda g:g.title)
  #parser.print_help()
  pa=parser.parse_args(sys.argv[1:])
  #print(pa,file=sys.stderr)
  if pa.length is not None:
    # We have to enforce our own check..
    if pa.offset is None or pa.filename is None:
      parser.error("length, offset and filename must all be supplied together")
 
  buf=bytearray(128*1024*1024)

  whole=not (pa.warc or pa.headers or pa.body)
  if pa.length is not None:
    process(pa,buf,pa.root,FPAT%list(pa.filename.split('/')),
            pa.offset,pa.length,whole)
    exit(0)
  if pa.index:
    CDX=regex.compile('length": "([0-9]*)", "offset": "([0-9]*)", "filename": "crawl-data/([^/]*)/segments/([^/]*)/warc/(.*\.gz)"')
    for l in sys.stdin:
      m=CDX.search(l)
      if m is None:
        print("index line problem: \"%s\""%l.lstrip(),file=sys.stderr)
        exit(2)
      f=FPAT%(m[3:6])
      process(pa,buf,pa.root,f,
              int(m[2]),int(m[1]),whole)
    exit(0)
if __name__ == "__main__":
    main()