#!/usr/bin/env python3
"""测试 Envoy 限流是否生效"""

import http.client
import json
import re

BASE = "127.0.0.1"
PORT = 18080

ENDPOINTS = [
    (
        "/api/v1/auth/login",
        10,
        {"username": "test", "password": "test", "phone": "0", "remember": True},
    ),
    (
        "/api/v1/auth/register",
        5,
        {"username": "test", "password": "test", "phone": "0"},
    ),
    ("/api/v1/auth/forgot-password/send", 3, {"email": "test@test.com"}),
    ("/api/v1/email/verification-code/send", 3, {"email": "test@test.com"}),
]


def get_cookies():
    conn = http.client.HTTPConnection(BASE, PORT)
    conn.request("GET", "/healthz")
    resp = conn.getresponse()
    resp.read()
    cookies = {}
    for h, v in resp.getheaders():
        if h.lower() == "set-cookie":
            m = re.match(r"([^=]+)=([^;]+)", v)
            if m:
                cookies[m.group(1)] = m.group(2)
    conn.close()
    return cookies


def post(cookies, path, body_dict):
    conn = http.client.HTTPConnection(BASE, PORT)
    token = cookies.get("__Host-XSRF-TOKEN", "")
    guard = cookies.get("__Host-XSRF-GUARD", "")
    headers = {
        "Content-Type": "application/json",
        "Cookie": f"__Host-XSRF-TOKEN={token}; __Host-XSRF-GUARD={guard}",
        "xsrf-token": token,
    }
    conn.request("POST", path, body=json.dumps(body_dict), headers=headers)
    resp = conn.getresponse()
    resp.read()
    rl_remain = resp.getheader("x-ratelimit-remaining", "-")
    ratelimited = resp.getheader("x-envoy-ratelimited", "")
    conn.close()
    return resp.status, rl_remain, bool(ratelimited)


def test_endpoint(cookies, path, limit, body_dict):
    count = limit + 3
    print(f"\n=== {path} (限额 {limit}/min, 发 {count} 次) ===")
    first_429 = None
    for i in range(1, count + 1):
        status, remain, limited = post(cookies, path, body_dict)
        tag = " ← 429!" if limited else ""
        print(f"  #{i:02d}  status={status}  remaining={remain}{tag}")
        if limited and first_429 is None:
            first_429 = i
    if first_429 == limit + 1:
        print(f"  ✓ 第 {first_429} 次触发限流，符合预期")
    elif first_429:
        print(f"  ✗ 第 {first_429} 次触发限流，预期第 {limit + 1} 次")
    else:
        print(f"  ✗ 未触发限流！")


print("获取 CSRF Cookie...")
cookies = get_cookies()
print(f"  token = {cookies.get('__Host-XSRF-TOKEN', '?')[:40]}...")

for path, limit, body in ENDPOINTS:
    test_endpoint(cookies, path, limit, body)