#!/usr/bin/env python3
"""目标入口闸门：唯一去 app 的 HTTP 路径，负责"关入口 + 排空在途请求"。

为什么必须有一个独立闸门，而不是在被代理的 WeKnora 上排空：
v0.7.1 的 app 收到 SIGTERM 会先关监听、随后因 use of closed network connection
直接 exit(1)，在途请求被切断。所以"入口已关、在途已归零"这条证据只能在真正的
连接入口上产生，并且必须落在**目标主机自己的**文件上（治理侧的维护协调器会先删除
这个文件、发 SIGTERM、等容器退出、再读回；读到旧文件就是伪造证据）。

连接语义（两条都被真实业务踩过，不是理论问题）：
  1. 上游读**不设超时**：`create_connection(timeout=...)` 只对 connect 生效，连上后
     必须 `settimeout(None)`。否则一条等待模型输出超过这个秒数的空闲流会被闸门自己
     掐断 —— 那是"闸门把正常业务掐了"，不是"连接已经排空"。
  2. 半关闭按 TCP 语义单向传播：客户端写完请求后关自己的写端（shutdown(SHUT_WR)）时，
     只把 EOF 传到上游，**不能**顺手把上游→客户端方向也拆掉，否则响应会被截断。
     只有真正的错误（RST 等）才双向拆除，避免另一侧线程永久挂住。

SIGTERM 语义（顺序是证据的一部分）：
  1. 立即关闭监听 —— 新连接被拒 = 入口已关；
  2. 等在途连接自然结束（上限 GATE_DRAIN_TIMEOUT）；
  3. 写 drain.json（临时文件 + fsync + rename，0600）后 exit 0（drained）
     / exit 2（drain_timeout，此时绝不能当成排空成功）。
"""

import json
import os
import signal
import socket
import sys
import threading
import time

GATE_PORT = int(os.environ.get("GATE_PORT", "8080"))
UPSTREAM_HOST = os.environ.get("GATE_UPSTREAM_HOST", "app")
UPSTREAM_PORT = int(os.environ.get("GATE_UPSTREAM_PORT", "8080"))
DRAIN_FILE = os.environ.get("GATE_DRAIN_FILE", "/gate-state/drain.json")
DRAIN_TIMEOUT = float(os.environ.get("GATE_DRAIN_TIMEOUT", "45"))
CONNECT_TIMEOUT = float(os.environ.get("GATE_CONNECT_TIMEOUT", "10"))

active = 0
served = 0
lock = threading.Lock()
listener = None
stopping = threading.Event()


def write_state(payload):
    if not os.path.isabs(DRAIN_FILE):
        print("gate: GATE_DRAIN_FILE must be an absolute path", file=sys.stderr, flush=True)
        os._exit(3)
    payload = dict(payload, at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), pid=os.getpid())
    directory = os.path.dirname(DRAIN_FILE)
    os.makedirs(directory, exist_ok=True)
    tmp = DRAIN_FILE + ".tmp"
    fd = os.open(tmp, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w") as handle:
        json.dump(payload, handle)
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(tmp, DRAIN_FILE)


def pump(src, dst):
    """单向搬运并传播半关闭。正常 EOF 只关对端写端；异常则双向拆除。"""
    clean = False
    try:
        while True:
            data = src.recv(65536)
            if not data:
                clean = True
                break
            dst.sendall(data)
    except OSError:
        clean = False
    finally:
        if clean:
            try:
                dst.shutdown(socket.SHUT_WR)
            except OSError:
                pass
        else:
            for sock in (src, dst):
                try:
                    sock.shutdown(socket.SHUT_RDWR)
                except OSError:
                    pass


def handle(client):
    global active, served
    with lock:
        if stopping.is_set():
            client.close()
            return
        active += 1
        served += 1
        current = active
    print(f"conn:open active={current} served={served}", flush=True)
    upstream = None
    try:
        upstream = socket.create_connection((UPSTREAM_HOST, UPSTREAM_PORT), timeout=CONNECT_TIMEOUT)
        # 建连超时只对 connect 生效；读超时必须清掉，否则长静默流会被掐断。
        upstream.settimeout(None)
        client.settimeout(None)
    except OSError:
        with lock:
            active -= 1
            remaining = active
        print(f"conn:close active={remaining} upstream-unreachable", flush=True)
        if upstream is not None:
            upstream.close()
        client.close()
        return
    left = threading.Thread(target=pump, args=(client, upstream), daemon=True)
    right = threading.Thread(target=pump, args=(upstream, client), daemon=True)
    left.start()
    right.start()

    def finish():
        global active
        try:
            left.join()
            right.join()
        finally:
            # 两个方向都结束才算这条连接离开：只看一个方向会把半关的连接当"已排空"。
            for sock in (client, upstream):
                try:
                    sock.close()
                except OSError:
                    pass
            with lock:
                active -= 1
                remaining = active
            print(f"conn:close active={remaining}", flush=True)

    threading.Thread(target=finish, daemon=True).start()


def on_term(_signum, _frame):
    stopping.set()
    if listener is not None:
        try:
            listener.close()
        except OSError:
            pass


def main():
    global listener
    signal.signal(signal.SIGTERM, on_term)
    signal.signal(signal.SIGINT, on_term)
    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind(("0.0.0.0", GATE_PORT))
    listener.listen(256)
    write_state({"phase": "serving", "active": 0, "served": 0})
    started = time.time()
    print(f"gate: listening on {GATE_PORT} -> {UPSTREAM_HOST}:{UPSTREAM_PORT} drain={DRAIN_FILE}",
          flush=True)
    while not stopping.is_set():
        try:
            client, _ = listener.accept()
        except OSError:
            break
        handle(client)

    with lock:
        active_at_close = active
    drain_started = time.time()
    write_state({"phase": "draining", "active": active_at_close, "served": served})
    deadline = drain_started + DRAIN_TIMEOUT
    while time.time() < deadline:
        with lock:
            if active == 0:
                break
        time.sleep(0.05)
    with lock:
        remaining = active
    drain_ms = int((time.time() - drain_started) * 1000)
    uptime_ms = int((time.time() - started) * 1000)
    if remaining != 0:
        write_state({"phase": "drain_timeout", "active": remaining, "served": served,
                     "active_at_close": active_at_close, "drain_ms": drain_ms,
                     "uptime_ms": uptime_ms})
        os._exit(2)
    write_state({"phase": "drained", "active": 0, "served": served,
                 "active_at_close": active_at_close, "drain_ms": drain_ms,
                 "uptime_ms": uptime_ms})
    os._exit(0)


if __name__ == "__main__":
    main()