#!/usr/bin/env python3
"""gate.py 的真实回归（无外部依赖、无副作用，只在本机 loopback 上跑）。

覆盖的是"闸门会不会掐断正常业务"和"排空证据是不是真的"这两件事，而不是 /health 200：

  1. 长静默流：上游发一半后静默 12 秒再发剩下的 —— 闸门必须照原样送达
     （旧实现 create_connection(timeout=10) 未清除读超时，会在这里失败）。
  2. 半关闭：客户端写完请求即 shutdown(SHUT_WR)，上游收到 EOF 后才应答 ——
     响应必须完整回来，且耗时证明上游确实等到了 EOF。
  3. 大包完整性：2MB 正文穿过闸门逐字节一致。
  4. 排空：在途连接存在时 SIGTERM —— 新连接被拒、在途连接照常收完、退出码 0、
     drain.json phase=drained 且 active_at_close>=1。
  5. 排空超时：上游永不应答、在途连接不结束 —— 退出码 2、phase=drain_timeout（不能算排空成功）。
  6. 上游不可达：连接被关掉而不是挂住，闸门本身仍存活并能正常排空。

用法：python3 deploy/cloud-test/test_gate.py   （打印 PASS/FAIL，任一 FAIL 退出码 1）
"""

import json
import os
import signal
import socket
import subprocess
import sys
import tempfile
import threading
import time
import traceback

HERE = os.path.dirname(os.path.abspath(__file__))
GATE = os.path.join(HERE, "gate.py")
SILENT_GAP = 12.0

failures = []


def check(name, condition, detail=""):
    if condition:
        print(f"PASS   {name} {detail}".rstrip(), flush=True)
    else:
        print(f"FAIL   {name} {detail}".rstrip(), flush=True)
        failures.append(name)


def free_port():
    sock = socket.socket()
    sock.bind(("127.0.0.1", 0))
    port = sock.getsockname()[1]
    sock.close()
    return port


class FakeUpstream:
    """按请求行决定行为的极简上游。"""

    def __init__(self):
        self.port = free_port()
        self.server = socket.socket()
        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server.bind(("127.0.0.1", self.port))
        self.server.listen(16)
        self.stopped = threading.Event()
        threading.Thread(target=self._serve, daemon=True).start()

    def _serve(self):
        while not self.stopped.is_set():
            try:
                client, _ = self.server.accept()
            except OSError:
                return
            threading.Thread(target=self._handle, args=(client,), daemon=True).start()

    def _handle(self, client):
        try:
            client.settimeout(None)
            data = b""
            while b"\r\n\r\n" not in data:
                chunk = client.recv(4096)
                if not chunk:
                    return
                data += chunk
            head, _, rest = data.partition(b"\r\n\r\n")
            path = head.split(b" ")[1].decode() if b" " in head else "/"
            if path == "/slow":
                client.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: 10\r\n\r\nabc")
                time.sleep(SILENT_GAP)
                client.sendall(b"defghij")
            elif path == "/half":
                # 等客户端半关闭（EOF）再应答：证明闸门把 EOF 单向传了过来。
                while True:
                    chunk = client.recv(4096)
                    if not chunk:
                        break
                time.sleep(0.8)
                body = b"after-eof"
                client.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n" % len(body) + body)
            elif path == "/echo":
                length = 0
                for line in head.split(b"\r\n"):
                    if line.lower().startswith(b"content-length:"):
                        length = int(line.split(b":")[1].strip())
                body = rest
                while len(body) < length:
                    chunk = client.recv(65536)
                    if not chunk:
                        break
                    body += chunk
                client.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n" % len(body) + body)
            elif path == "/hold":
                client.sendall(b"HTTP/1.1 200 OK\r\n\r\n")
                time.sleep(3600)
            else:
                client.sendall(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n")
        except OSError:
            pass
        finally:
            try:
                client.close()
            except OSError:
                pass

    def close(self):
        self.stopped.set()
        try:
            self.server.close()
        except OSError:
            pass


class Gate:
    def __init__(self, upstream_port, drain_timeout=45.0):
        self.port = free_port()
        self.dir = tempfile.mkdtemp(prefix="gate-test-")
        self.drain = os.path.join(self.dir, "drain.json")
        env = dict(os.environ)
        env.update({
            "GATE_PORT": str(self.port),
            "GATE_UPSTREAM_HOST": "127.0.0.1",
            "GATE_UPSTREAM_PORT": str(upstream_port),
            "GATE_DRAIN_FILE": self.drain,
            "GATE_DRAIN_TIMEOUT": str(drain_timeout),
        })
        self.process = subprocess.Popen([sys.executable, GATE, ], env=env,
                                        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        deadline = time.time() + 10
        while time.time() < deadline:
            try:
                socket.create_connection(("127.0.0.1", self.port), timeout=0.2).close()
                return
            except OSError:
                if self.process.poll() is not None:
                    raise RuntimeError("gate exited early: " + (self.process.stdout.read() or ""))
                time.sleep(0.05)
        raise RuntimeError("gate did not start listening")

    def connect(self):
        return socket.create_connection(("127.0.0.1", self.port), timeout=30)

    def state(self):
        with open(self.drain) as handle:
            return json.load(handle)

    def terminate(self, timeout=20):
        self.process.send_signal(signal.SIGTERM)
        try:
            self.process.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            self.process.kill()
            self.process.wait(5)
            return None
        return self.process.returncode

    def logs(self):
        try:
            return self.process.stdout.read() or ""
        except Exception:
            return ""


def read_all(conn, timeout=30):
    conn.settimeout(timeout)
    chunks = []
    try:
        while True:
            data = conn.recv(65536)
            if not data:
                break
            chunks.append(data)
    except socket.timeout:
        chunks.append(b"<TIMEOUT>")
    except ConnectionResetError:
        # 对端 RST（例如上游不可达时闸门立即拆连接）也是"连接已结束"，不是挂住。
        pass
    return b"".join(chunks)


def test_long_silent_stream():
    upstream = FakeUpstream()
    gate = Gate(upstream.port)
    try:
        conn = gate.connect()
        conn.sendall(b"GET /slow HTTP/1.1\r\nHost: gate\r\n\r\n")
        started = time.time()
        payload = read_all(conn, timeout=30)
        elapsed = time.time() - started
        conn.close()
        check("long-silent-stream/complete", b"abcdefghij" in payload,
              f"({elapsed:.1f}s, {len(payload)}B)")
        check("long-silent-stream/survived-gap", elapsed >= SILENT_GAP - 1,
              f"gap={SILENT_GAP}s elapsed={elapsed:.1f}s")
    finally:
        gate.terminate()
        upstream.close()


def test_half_close():
    upstream = FakeUpstream()
    gate = Gate(upstream.port)
    try:
        conn = gate.connect()
        conn.sendall(b"GET /half HTTP/1.1\r\nHost: gate\r\n\r\n")
        conn.shutdown(socket.SHUT_WR)   # 客户端半关闭：只关自己的写端
        started = time.time()
        payload = read_all(conn, timeout=15)
        elapsed = time.time() - started
        conn.close()
        check("half-close/response-delivered", b"after-eof" in payload, f"({elapsed:.1f}s)")
        check("half-close/eof-propagated", elapsed >= 0.5, f"elapsed={elapsed:.1f}s")
    finally:
        gate.terminate()
        upstream.close()


def test_large_body_integrity():
    upstream = FakeUpstream()
    gate = Gate(upstream.port)
    try:
        body = (b"0123456789abcdef" * 131072)  # 2 MiB
        conn = gate.connect()
        conn.sendall(b"POST /echo HTTP/1.1\r\nHost: gate\r\nContent-Length: %d\r\n\r\n" % len(body))
        conn.sendall(body)
        payload = read_all(conn, timeout=30)
        conn.close()
        head, _, received = payload.partition(b"\r\n\r\n")
        check("large-body/integrity", received == body,
              f"sent={len(body)}B got={len(received)}B head={head.splitlines()[:1]}")
    finally:
        gate.terminate()
        upstream.close()


def test_drain_with_inflight():
    upstream = FakeUpstream()
    gate = Gate(upstream.port, drain_timeout=45)
    try:
        conn = gate.connect()
        conn.sendall(b"GET /slow HTTP/1.1\r\nHost: gate\r\n\r\n")
        time.sleep(0.5)                       # 确保在途连接已被计数
        sent_at = time.time()
        gate.process.send_signal(signal.SIGTERM)
        time.sleep(0.5)
        refused = False
        try:
            late = socket.create_connection(("127.0.0.1", gate.port), timeout=2)
            late.settimeout(2)
            try:
                late.sendall(b"GET /slow HTTP/1.1\r\nHost: gate\r\n\r\n")
                refused = late.recv(1) == b""
            except OSError:
                refused = True
            late.close()
        except OSError:
            refused = True
        check("drain/new-connections-refused", refused)
        payload = read_all(conn, timeout=40)
        conn.close()
        check("drain/inflight-completed", b"abcdefghij" in payload, f"{len(payload)}B")
        code = gate.terminate(timeout=30)
        check("drain/exit-zero", code == 0, f"exit={code}")
        state = gate.state()
        fresh = os.path.getmtime(gate.drain) >= sent_at - 1
        check("drain/evidence", state.get("phase") == "drained" and state.get("active_at_close", 0) >= 1,
              json.dumps({k: state.get(k) for k in ("phase", "active_at_close", "served", "drain_ms")}))
        check("drain/evidence-fresh", fresh, f"mtime ok={fresh}")
    finally:
        if gate.process.poll() is None:
            gate.terminate()
        upstream.close()


def test_drain_timeout_is_failure():
    upstream = FakeUpstream()
    gate = Gate(upstream.port, drain_timeout=1)
    try:
        conn = gate.connect()
        conn.sendall(b"GET /hold HTTP/1.1\r\nHost: gate\r\n\r\n")
        time.sleep(0.5)
        code = gate.terminate(timeout=15)
        check("drain-timeout/exit-nonzero", code == 2, f"exit={code}")
        state = gate.state()
        check("drain-timeout/phase", state.get("phase") == "drain_timeout",
              json.dumps({k: state.get(k) for k in ("phase", "active", "active_at_close")}))
        conn.close()
    finally:
        if gate.process.poll() is None:
            gate.terminate()
        upstream.close()


def test_upstream_unreachable():
    dead = free_port()
    gate = Gate(dead)
    try:
        conn = gate.connect()
        conn.sendall(b"GET / HTTP/1.1\r\nHost: gate\r\n\r\n")
        started = time.time()
        payload = read_all(conn, timeout=5)
        elapsed = time.time() - started
        conn.close()
        check("upstream-dead/connection-closed", b"<TIMEOUT>" not in payload and elapsed < 4,
              f"{elapsed:.1f}s {len(payload)}B")
        # 闸门本身还活着：仍能建连（会在上游阶段失败，但不挂住）。
        code = gate.terminate(timeout=10)
        check("upstream-dead/gate-drains", code == 0, f"exit={code}")
    finally:
        if gate.process.poll() is None:
            gate.terminate()


def main():
    for test in (
        test_long_silent_stream,
        test_half_close,
        test_large_body_integrity,
        test_drain_with_inflight,
        test_drain_timeout_is_failure,
        test_upstream_unreachable,
    ):
        try:
            test()
        except Exception:
            traceback.print_exc()
            failures.append(test.__name__)
    print()
    print(f"summary: {len(failures)} failure(s)" + (":" + ",".join(failures) if failures else ""))
    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main())