#!/usr/bin/env python3
"""Small live two-origin failover and active-exchange regression, not a benchmark."""
from __future__ import annotations

import argparse
from collections import Counter
import copy
import hashlib
from http.client import HTTPConnection
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
from pathlib import Path
import threading
import time
import traceback

from common import Fixture, http, seal_results, write_json


class Origin:
    def __init__(self, port, name):
        self.name, self.port = name, port
        self.fail_get = False
        self.counts, self.post_hashes = Counter(), []
        self.release, self.lock = threading.Event(), threading.Lock()
        origin = self

        class Handler(BaseHTTPRequestHandler):
            protocol_version = "HTTP/1.1"

            def log_message(self, *_):
                pass

            def do_GET(self):
                self.respond()

            def do_POST(self):
                self.respond()

            def respond(self):
                body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
                with origin.lock:
                    origin.counts[(self.command, self.path)] += 1
                    if self.command == "POST":
                        origin.post_hashes.append(hashlib.sha256(body).hexdigest())
                hold = self.path == "/hold"
                status = 503 if self.command == "POST" or origin.fail_get else 200
                payload = b"firstlast" if hold else origin.name.encode()
                try:
                    self.send_response(status)
                    self.send_header("Content-Length", str(len(payload)))
                    self.send_header("Content-Type", "application/octet-stream")
                    self.send_header("Cache-Control", "no-store")
                    # Fresh connections make refused-listener tests independent
                    # of already accepted upstream keepalive connections.
                    self.send_header("Connection", "close")
                    self.end_headers()
                    if hold:
                        self.wfile.write(payload[:5])
                        self.wfile.flush()
                        origin.release.wait(15)
                        self.wfile.write(payload[5:])
                    else:
                        self.wfile.write(payload)
                    self.wfile.flush()
                except (BrokenPipeError, ConnectionResetError):
                    pass
                self.close_connection = True

        self.server = ThreadingHTTPServer(("127.0.0.1", port), Handler)
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
        self.thread.start()
        self.closed = False

    def close(self):
        self.release.set()
        if not self.closed:
            self.server.shutdown()
            self.server.server_close()
            self.thread.join(timeout=3)
            self.closed = True


def backend_views(rust):
    status, _, body = rust.admin("/api/status")
    if status != 200:
        raise RuntimeError(f"status unavailable: {status}")
    found = {}

    def walk(value):
        if isinstance(value, dict):
            if value.get("id") in ("a", "b") and "active" in value and "status" in value:
                found[value["id"]] = {key: value[key] for key in value
                    if key in ("id", "status", "active", "requests", "errors", "bytes_out")
                    or key.startswith("retry_") or key.startswith("connect_errors")}
            for child in value.values():
                walk(child)
        elif isinstance(value, list):
            for child in value:
                walk(child)
    walk(json.loads(body))
    return found


def run(check, fixture, binary, label):
    a, b = Origin(fixture.origin_port, "a"), Origin(fixture.port(), "b")
    try:
        rust = fixture.rust(binary, label)
        rust.login()
        config = copy.deepcopy(rust.base)
        vhost = config["vhosts"][0]
        vhost["backends"] = [{"id": origin.name, "url": f"http://127.0.0.1:{origin.port}",
                               "allow_private": True, "preserve_host": False} for origin in [a, b]]
        vhost["load_balancing"] = {"algorithm": "round_robin", "health_check": {"mode": "off"},
            "slow_start_seconds": 0, "passive_health": {"enabled": False},
            "retry": {"max_retries": 1, "retry_statuses": [503]}}
        vhost["routes"] = [{"id": "hold-a", "path_prefix": "/hold", "backend_id": "a"}]
        rust.apply(config, "lb-two-origins")
        replies = [http(rust.reverse_port, "/rr") for _ in range(8)]
        counts = Counter(body.decode() for status, _, body in replies if status == 200)
        check("lb.round_robin_both_origins", counts == {"a": 4, "b": 4}, counts=dict(counts))

        a.fail_get = True
        replies = [http(rust.reverse_port, "/retry-status") for _ in range(6)]
        check("lb.get_503_failover", all(status == 200 and body == b"b" for status, _, body in replies)
              and a.counts[("GET", "/retry-status")] > 0,
              failed_origin_attempts=a.counts[("GET", "/retry-status")])
        views = backend_views(rust)
        check("lb.retry_observations_recorded", sum(view.get("retry_attempts_total", 0) for view in views.values()) > 0,
              backends=views)
        a.fail_get = False
        payload = b"one POST must be attempted once"
        replies = [http(rust.reverse_port, "/post-no-replay", "POST", body=payload) for _ in range(4)]
        attempts = sum(origin.counts[("POST", "/post-no-replay")] for origin in [a, b])
        check("lb.post_503_never_replayed", all(status == 503 for status, _, _ in replies) and attempts == 4,
              client_requests=4, origin_attempts=attempts)
        check("lb.post_body_preserved", len(a.post_hashes + b.post_hashes) == 4 and
              all(digest == hashlib.sha256(payload).hexdigest() for digest in a.post_hashes + b.post_hashes))

        for mode in ["drain", "disable"]:
            vhost["backends"][0].update({"enabled": True, "drain": False})
            rust.apply(config, f"lb-{mode}-before")
            a.release.clear()
            connection = HTTPConnection("127.0.0.1", rust.reverse_port, timeout=10)
            try:
                connection.request("GET", "/hold")
                response = connection.getresponse()
                first = response.read(5)
                before = backend_views(rust)
                check(f"lb.{mode}.lease_spans_response_body", response.status == 200 and first == b"first"
                      and before.get("a", {}).get("active") == 1, backends=before)
                vhost["backends"][0].update({"drain": mode == "drain", "enabled": mode != "disable"})
                rust.apply(config, f"lb-{mode}-active")
                replies = [http(rust.reverse_port, "/new") for _ in range(4)]
                check(f"lb.{mode}.new_requests_use_healthy_member",
                      all(status == 200 and body == b"b" for status, _, body in replies))
                during = backend_views(rust)
                check(f"lb.{mode}.active_lease_survives_hot_edit", during.get("a", {}).get("active") == 1,
                      backends=during)
                a.release.set()
                last = response.read()
                check(f"lb.{mode}.admitted_body_finishes", first + last == b"firstlast")
                for _ in range(20):
                    after = backend_views(rust)
                    if after.get("a", {}).get("active") == 0:
                        break
                    time.sleep(.05)
                check(f"lb.{mode}.lease_released_after_body", after.get("a", {}).get("active") == 0,
                      backends=after)
            finally:
                a.release.set()
                connection.close()

        vhost["backends"][0].update({"enabled": True, "drain": False})
        rust.apply(config, "lb-connect-refused")
        a.close()
        replies = [http(rust.reverse_port, "/retry-connect") for _ in range(6)]
        views = backend_views(rust)
        check("lb.connect_refused_failover", all(status == 200 and body == b"b" for status, _, body in replies)
              and views.get("a", {}).get("connect_errors_total", 0) > 0, backends=views)
        rust.reset()
        check("lb.static_vhost_survives_reload", http(rust.static_port)[0] == 200)
    finally:
        a.close()
        b.close()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--binary", action="append", required=True, metavar="LABEL=/ABS/PATH")
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    if args.out.exists() and any(args.out.iterdir()):
        parser.error("--out must be new or empty")
    reports = {}
    for specification in args.binary:
        label, binary = specification.split("=", 1)
        cases = []

        def check(case, passed, **details):
            cases.append({"case": case, "passed": passed, **details})
            print(f"{label} {'PASS' if passed else 'FAIL'} {case}", flush=True)
            write_json(args.out / label / "checks.json", cases)

        try:
            with Fixture(args.out / label) as fixture:
                run(check, fixture, Path(binary), label)
        except Exception as exc:
            check("lb.runner", False, error=str(exc), traceback=traceback.format_exc())
        report = {"passed": sum(case["passed"] for case in cases),
                  "failed": sum(not case["passed"] for case in cases), "cases": cases,
                  "simulation": False, "kind": "live RustProxy two-origin HTTP/1.1 regression",
                  "not_covered": ["sticky affinity", "active health probes", "long-running connection soak",
                                  "TLS origin failures", "DNS rotation"]}
        reports[label] = report
        write_json(args.out / label / "report.json", report)
    write_json(args.out / "comparison.json", reports)
    seal_results(args.out)
    return 1 if any(report["failed"] for report in reports.values()) else 0


if __name__ == "__main__":
    raise SystemExit(main())