#! /usr/bin/python3

# TODO move into a real web server?
# TODO allow different teams to coexist?
# TODO wait till after a probe confirms it has finished before deleting a target
# file? what happens when probes check out a file and don't complete it?

from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import os
import queue
import re
import shutil
import sys
import time

from ark_http_handler import ArkHTTPServer, ArkHTTPRequestHandler

# run a proxy in front of it to check ssl client certificates
HOSTNAME = "127.0.0.1"
PORT = 8081
MINIMUM_SECONDS = 300
DEFAULT_DELAY = 300
DISABLED_DELAY = 60 * 60 * 24
TARGET_FILE_REGEX = re.compile("targets\\.c(?P<cycle>\\d+)\\.\\d+$")


def _usage():
    print("usage: ark-team-probing-targets-server $dir\n"
          "       dir:   directory to read target files from\n")


class TeamProbingHandler(ArkHTTPRequestHandler):
    def do_GET(self):
        # assume each client has a unique address
        #client = self.address_string()
        client = self.get_client_address()

        # don't serve nodes that shouldn't be running team-probing
        name = self.get_name_from_certificate()
        monitor = self.server.monitors.get(name)
        if monitor is None:
            self.send_retry(HTTPStatus.SERVICE_UNAVAILABLE, DISABLED_DELAY)
            print(f"{client}/{name}: client unknown", file=sys.stderr)
            return

        if "IPv4 team probing" not in monitor.get("activities", []):
            self.send_retry(HTTPStatus.SERVICE_UNAVAILABLE, DISABLED_DELAY)
            print(f"{client}/{name}: team-probing disabled", file=sys.stderr)
            return

        # processing a target file should take ~8-10 minutes, so check that
        # the client isn't retrying too quickly and wasting target files
        delay = self.should_retry()
        if delay > 0:
            self.send_retry(HTTPStatus.TOO_MANY_REQUESTS, delay)
            print(f"{client}/{name}: too soon, retry in {delay}s", file=sys.stderr)
            return

        # get the next available target file
        filename = self.get_next_target_file()
        if filename is None:
            # no more target files available, ask the client to retry later
            self.send_retry(HTTPStatus.SERVICE_UNAVAILABLE)
            print(f"{client}/{name}: queue empty, retry in {DEFAULT_DELAY}s",
                    file=sys.stderr)
            # should wait for all open connections to finish, and reload
            # a new filelist without worrying about thread interactions
            self.server.shutdown()
            return

        # make sure the cycle id is encoded in the filename
        match = TARGET_FILE_REGEX.search(filename)
        if match is None:
            self.send_retry(HTTPStatus.SERVICE_UNAVAILABLE)
            print(f"{client}/{name}: bad file, retry in {DEFAULT_DELAY}s",
                    file=sys.stderr)
            return

        # hopefully we can open the file
        try:
            with open(filename, "rb") as targetfile:
                # send a new target file to the client
                self.send_response(HTTPStatus.OK)
                self.send_header("Content-Type", "text/plain")
                self.send_header("X-ark-team-probing-cycle",
                        match.group("cycle"))
                self.end_headers()
                print(f"sending targets file {filename} to {client}/{name}",
                        file=sys.stderr)
                shutil.copyfileobj(targetfile, self.wfile)
        except IOError:
            self.send_retry(HTTPStatus.SERVICE_UNAVAILABLE)
            print(f"{client}/{name}: bad file, retry in {DEFAULT_DELAY}s",
                    file=sys.stderr)
            return
        os.unlink(filename)

    def get_next_target_file(self):
        try:
            filename = self.server.targets.get_nowait()
        except queue.Empty:
            return None
        return filename

    def send_retry(self, status, after=DEFAULT_DELAY):
        self.send_response(status)
        self.send_header("Retry-After", after)
        self.end_headers()

    def should_retry(self):
        # assume each client has a unique address
        client = self.get_client_address()
        if client in self.server.limits:
            elapsed = int(time.time()) - self.server.limits[client]
            if elapsed < MINIMUM_SECONDS:
                return MINIMUM_SECONDS - elapsed
        # record the most recent allowed request time
        self.server.limits[client] = int(time.time())
        return 0


def load_target_files(target_queue, directory):
    if not target_queue.empty():
        return
    print(f"loading target files from {directory}")
    for target_file in sorted(os.listdir(directory)):
        if TARGET_FILE_REGEX.search(target_file):
            target_queue.put(os.path.join(directory, target_file))
    print(f"loaded {target_queue.qsize()} target files from {directory}")


if __name__ == "__main__":
    # TODO who is responsible for cycle number?
    # TODO if this exits when out of targets, cycle could be a command line arg
    # though that might not work well if run properly as wsgi
    if len(sys.argv) != 2 or not os.path.isdir(sys.argv[1]):
        _usage()
        sys.exit(1)

    server = ArkHTTPServer((HOSTNAME, PORT), TeamProbingHandler)
    server.allow_reuse_address = True
    server.request_queue_size = 0
    server.limits = {}
    server.targets = queue.Queue()

    while True:
        # load all the target files from the given directory
        load_target_files(server.targets, sys.argv[1])

        # If shutdown() is called and serve_forever() leaves the loop, then
        # any future calls to shutdown() set the shutdown_request flag again.
        # So when we call serve_forever() again after reloading target files
        # it immediately exits the loop and tries to reload again.
        # TODO Could wrap server.handle_request() in a loop with a public
        # condition variable, but for for now we can reach inside and reset
        # shutdown_request.
        # TODO alternatively, let the program exit and get restarted
        server._BaseServer__shutdown_request = False

        print(f"Serving target files on {HOSTNAME}:{PORT}")
        try:
            server.serve_forever()
        except KeyboardInterrupt:
            server.server_close()
            sys.exit(0)

        # brief wait to hopefully catch any more shutdown() calls
        time.sleep(2)
