#!/usr/bin/python3
import http.server
import socket
import subprocess
import tempfile
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        data = b"chisel-autopkgtest-forwarded\n"
        self.send_response(200)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def log_message(self, *args):
        pass

def free_port():
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        return sock.getsockname()[1]

with tempfile.TemporaryDirectory() as directory:
    processes = []
    handles = []
    target = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
    threading.Thread(target=target.serve_forever, daemon=True).start()
    server_port, forward_port = free_port(), free_port()
    while forward_port == server_port:
        forward_port = free_port()
    commands = [
        ["chisel", "server", "--host", "127.0.0.1", "--port", str(server_port)],
        ["chisel", "client", "--max-retry-count", "10",
         f"http://127.0.0.1:{server_port}",
         f"127.0.0.1:{forward_port}:127.0.0.1:{target.server_port}"],
    ]
    try:
        for index, command in enumerate(commands):
            log = open(Path(directory) / f"process-{index}.log", "w")
            handles.append(log)
            processes.append(subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT))
        # Do not allow proxy environment variables to redirect the loopback test.
        opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
        deadline = time.monotonic() + 30
        while True:
            for process in processes:
                assert process.poll() is None, "Chisel exited before forwarding a request"
            try:
                with opener.open(f"http://127.0.0.1:{forward_port}/", timeout=2) as response:
                    assert response.status == 200
                    assert response.read() == b"chisel-autopkgtest-forwarded\n"
                break
            except (urllib.error.URLError, ConnectionError, TimeoutError):
                if time.monotonic() >= deadline:
                    raise
                time.sleep(0.1)
    except Exception:
        for path in Path(directory).glob("*.log"):
            print(path.read_text())
        raise
    finally:
        for process in reversed(processes):
            process.terminate()
        for process in processes:
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait()
        for log in handles:
            log.close()
        target.shutdown()
        target.server_close()
