#!/usr/bin/env python3

import http.server
import hashlib
import urllib.request
import http.server
import subprocess
import sys
import signal
import itertools
import requests


class CachingHandler(http.server.BaseHTTPRequestHandler):
    def _cached_read(self, input, chunksize=4096):
        # >>> datetime.datetime.utcnow().strftime("%a, %d %b %Y %H:%M:%S GMT")
        with open(input, 'rb') as data:
            while True:
                chunk = data.read(chunksize)
                if not chunk:
                    break

                yield chunk

    def _saving_read(self, output):
        with requests.get(self.path, stream=True) as src:
            if src.status_code != 200:
                self.send_error(src.status_code)
                return

            self.send_response(200)
            self.end_headers()

            try:
                with open(output, 'wb') as data:
                    for chunk in src.iter_content(chunk_size=4096):
                        if chunk:
                            data.write(chunk)
                            yield(chunk)

            except Exception as err:
                os.unlink(output)
                raise err

    def do_GET(self):
        h = hashlib.sha256()
        h.update(self.path.encode())
        path = os.path.join(args.store, h.hexdigest() + ".data")

        chunker = self._saving_read if not os.path.exists(path) else self._cached_read
        for chunk in chunker(path):
            self.wfile.write(chunk)


if __name__ == '__main__':
    import argparse
    import os

    parser = argparse.ArgumentParser()
    parser.add_argument('--store', type=str, required=True)
    parser.add_argument('--address', type=str, default="0.0.0.0")
    parser.add_argument('--port', type=int, default=8000)

    parser.add_argument('cmd', type=str, nargs=1)
    parser.add_argument('args', type=str, nargs=argparse.REMAINDER)
    args = parser.parse_args()

    # Instantiate the server before we run the child process
    os.makedirs(args.store, exist_ok=True)

    server = http.server.HTTPServer(
        (args.address, args.port),
        CachingHandler
    )

    pid = os.fork()
    if pid == 0:
        # Start the child process. Tell the parent to shutdown the server when
        # we are finished.
        res = subprocess.run(
            ['env', f"http_proxy={args.address}:{args.port}"] + args.cmd + args.args,
            stderr=sys.stderr, stdout=sys.stdout
        )

        os.kill(os.getppid(), signal.SIGHUP)
        sys.exit(res.returncode)
    else:
        # Run a server until we get a SIGHUP, then wait for the child to exit.
        server.serve_forever()
        (_, status) = os.waitpid(pid)
        sys.exit(status)