"""Same-language Python stress / benchmark harness.

고성능 병렬 운용 기준선 Milestone의 lang-baseline Task에서, TypeScript stress.ts와 같은 출력
계약(ROW|/SKIP|/SUMMARY|)으로 Python same-language TCP/WS baseline을 남긴다. 절대 성능 합격선은
고정하지 않고 local 환경 baseline(throughput, p50/p95/p99 latency, peak RSS)을 기록하며, 안정성
합격선만 hard fail로 둔다: timeout 0, nonce mismatch 0, response type mismatch 0, connection별
FIFO 위반 0, 종료 후 pending leak 0.

실행: python3 bench/stress.py [--mode quick|full] [--transport tcp,ws] [--profiles a,b]
"""

from __future__ import annotations

import argparse
import asyncio
import socket
import sys
import time
from contextlib import suppress
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from proto_socket.communicator import type_name_of
from proto_socket.packets.message_common_pb2 import TestData
from proto_socket.tcp_client import connect_tcp
from proto_socket.tcp_server import TcpServer
from proto_socket.ws_client import connect_ws
from proto_socket.ws_server import WsServer

HOST = "127.0.0.1"
LANGUAGE = "Python"
WS_PATH = "/"
ALL_PROFILES = ["roundtrip", "burst", "sustained", "parallel", "slow-mix", "payload"]

# payload size matrix 축. quick은 작은 샘플(1KB/64KB), full은 Milestone 기준(1KB/64KB/1MB)을 측정한다.
PAYLOAD_SEED = "0123456789abcdef"
PAYLOAD_SIZES_QUICK = [1024, 65536]
PAYLOAD_SIZES_FULL = [1024, 65536, 1048576]

_rows_emitted = 0
_total_violations = 0


class Stability:
    """안정성 위반 카운터. 0이 아니면 해당 row는 FAIL이고 프로세스도 non-zero exit한다."""

    __slots__ = ("timeouts", "nonce_mismatch", "type_mismatch", "fifo_violations", "pending_leak")

    def __init__(self) -> None:
        self.timeouts = 0
        self.nonce_mismatch = 0
        self.type_mismatch = 0
        self.fifo_violations = 0
        self.pending_leak = 0

    def violations(self) -> int:
        return (
            self.timeouts
            + self.nonce_mismatch
            + self.type_mismatch
            + self.fifo_violations
            + self.pending_leak
        )


def parser_map():
    return {type_name_of(TestData): TestData.FromString}


def log(line: str) -> None:
    print(line, file=sys.stderr, flush=True)


def mem_mb() -> float:
    try:
        with open("/proc/self/statm", encoding="ascii") as handle:
            resident_pages = int(handle.read().split()[1])
        return resident_pages * 4096 / (1024 * 1024)
    except OSError:
        import resource

        return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024


def free_port() -> int:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        sock.bind((HOST, 0))
        return sock.getsockname()[1]
    finally:
        sock.close()


def test_data_payload_bytes(message: str) -> int:
    return len(TestData(index=0, message=message).SerializeToString())


def payload_filler(size: int) -> str:
    """size 바이트 길이의 deterministic filler string을 만든다."""
    if size <= 0:
        return ""
    reps = -(-size // len(PAYLOAD_SEED))
    return (PAYLOAD_SEED * reps)[:size]


def payload_label(num_bytes: int) -> str:
    """payload 크기를 사람이 읽는 축 이름으로 바꾼다(1024 -> 1KB, 1048576 -> 1MB)."""
    if num_bytes >= 1048576 and num_bytes % 1048576 == 0:
        return f"{num_bytes // 1048576}MB"
    if num_bytes >= 1024 and num_bytes % 1024 == 0:
        return f"{num_bytes // 1024}KB"
    return f"{num_bytes}B"


def percentile(sorted_values: list[float], p: float) -> float:
    if not sorted_values:
        return 0.0
    rank = max(0, min(len(sorted_values) - 1, -(-int(p * len(sorted_values)) // 100) - 1))
    return sorted_values[rank]


def classify_request_error(message: str, stability: Stability) -> None:
    if "timeout" in message:
        stability.timeouts += 1
    elif "type mismatch" in message:
        stability.type_mismatch += 1
    else:
        stability.nonce_mismatch += 1


def emit_row(
    *,
    profile: str,
    axis: str,
    transport: str,
    payload_bytes: int,
    client_count: int,
    requests: int,
    throughput: float,
    p50: float,
    p95: float,
    p99: float,
    mem: float,
    stability: Stability,
) -> None:
    global _rows_emitted, _total_violations
    violations = stability.violations()
    status = "PASS" if violations == 0 else "FAIL"
    _total_violations += violations
    _rows_emitted += 1
    fields = [
        "ROW",
        profile,
        axis,
        LANGUAGE,
        transport,
        str(payload_bytes),
        str(client_count),
        str(requests),
        f"{throughput:.1f}",
        f"{p50:.3f}",
        f"{p95:.3f}",
        f"{p99:.3f}",
        str(stability.timeouts),
        str(stability.nonce_mismatch),
        str(stability.type_mismatch),
        str(stability.fifo_violations),
        str(stability.pending_leak),
        "0",  # queueBacklog: same-language baseline에는 inbound gateway가 없다.
        "0",  # gatewayBacklog
        "off",
        f"{mem:.1f}",
        status,
    ]
    print("|".join(fields), flush=True)


def emit_skip(profile: str, axis: str, transport: str, reason: str) -> None:
    print(f"SKIP|{profile}|{axis}|{LANGUAGE}|{transport}|{reason}", flush=True)


def summarize(
    profile: str,
    axis: str,
    transport: str,
    latencies: list[float],
    elapsed_ms: float,
    client_count: int,
    stability: Stability,
    mem: float,
) -> None:
    ordered = sorted(latencies)
    throughput = (len(latencies) / elapsed_ms * 1000.0) if elapsed_ms > 0 else 0.0
    emit_row(
        profile=profile,
        axis=axis,
        transport=transport,
        payload_bytes=test_data_payload_bytes("req-0"),
        client_count=client_count,
        requests=len(latencies),
        throughput=throughput,
        p50=percentile(ordered, 50),
        p95=percentile(ordered, 95),
        p99=percentile(ordered, 99),
        mem=mem,
        stability=stability,
    )


def emit_throughput(
    profile: str,
    axis: str,
    transport: str,
    count: int,
    elapsed_ms: float,
    client_count: int,
    payload_bytes: int,
    stability: Stability,
    mem: float,
) -> None:
    throughput = (count / elapsed_ms * 1000.0) if elapsed_ms > 0 else 0.0
    emit_row(
        profile=profile,
        axis=axis,
        transport=transport,
        payload_bytes=payload_bytes,
        client_count=client_count,
        requests=count,
        throughput=throughput,
        p50=0.0,
        p95=0.0,
        p99=0.0,
        mem=mem,
        stability=stability,
    )


def make_server(transport: str, port: int, on_connected):
    if transport == "tcp":
        server = TcpServer(HOST, port, interval_sec=0, wait_sec=0, parser_map=parser_map())
    elif transport == "ws":
        server = WsServer(HOST, port, WS_PATH, interval_sec=0, wait_sec=0, parser_map=parser_map())
    else:
        raise ValueError(f"unknown transport {transport!r}")
    server.on_client_connected = on_connected
    return server


async def dial(transport: str, port: int):
    if transport == "tcp":
        return await connect_tcp(HOST, port, 0, 0, parser_map())
    if transport == "ws":
        return await connect_ws(HOST, port, WS_PATH, 0, 0, parser_map())
    raise ValueError(f"unknown transport {transport!r}")


async def dial_with_retry(transport: str, port: int):
    loop = asyncio.get_running_loop()
    deadline = loop.time() + 3.0
    last_error: Exception | None = None
    while loop.time() < deadline:
        try:
            return await asyncio.wait_for(dial(transport, port), 0.3)
        except Exception as exc:  # noqa: BLE001
            last_error = exc
            await asyncio.sleep(0.05)
    raise TimeoutError(f"connect {transport}:{port} timed out: {last_error}")


def make_echo_server(transport: str, port: int):
    # request-response 축에서는 client가 concurrency C로 동시에 send하므로 frame 도착 순서가
    # 결정적이지 않다. 따라서 여기서는 FIFO를 측정하지 않고 응답 정확성/timeout만 본다. per-connection
    # FIFO는 단일 connection 순차 send인 burst 축에서 검증한다.
    return make_server(
        transport,
        port,
        lambda client: client.communicator.add_request_listener(
            type_name_of(TestData),
            lambda req: TestData(index=req.index * 2, message="echo:" + req.message),
        ),
    )


def make_slow_mix_server(transport: str, port: int, slow_every: int, delay: float):
    async def on_request(req: TestData) -> TestData:
        if (req.index + 1) % slow_every == 0:
            await asyncio.sleep(delay)
        return TestData(index=req.index * 2, message="echo:" + req.message)

    return make_server(
        transport,
        port,
        lambda client: client.communicator.add_request_listener(type_name_of(TestData), on_request),
    )


async def run_request_load(client, total: int, concurrency: int, timeout: float, stability: Stability) -> list[float]:
    latencies: list[float] = []
    issued = 0
    next_index = 0

    async def one_request(index: int) -> None:
        started = time.perf_counter()
        try:
            res = await client.communicator.send_request(
                TestData(index=index, message=f"req-{index}"), TestData, timeout=timeout
            )
            latencies.append((time.perf_counter() - started) * 1000.0)
            if res.index != index * 2 or res.message != f"echo:req-{index}":
                stability.nonce_mismatch += 1
        except Exception as exc:  # noqa: BLE001
            classify_request_error(str(exc), stability)

    while issued < total:
        batch_size = min(concurrency, total - issued)
        indices = list(range(next_index, next_index + batch_size))
        next_index += batch_size
        issued += batch_size
        await asyncio.gather(*(one_request(i) for i in indices))

    return latencies


async def profile_roundtrip(transport: str, mode: str) -> None:
    concurrencies = [1, 16, 64, 256]
    batches = 2 if mode == "quick" else 20
    timeout = 5.0 if mode == "quick" else 15.0
    log(f"[roundtrip] transport={transport} mode={mode} concurrencies=1,16,64,256 batches/level={batches}")
    for concurrency in concurrencies:
        stability = Stability()
        total = concurrency * batches
        port = free_port()
        server = make_echo_server(transport, port)
        await server.start()
        try:
            client = await dial_with_retry(transport, port)
            try:
                started = time.perf_counter()
                latencies = await run_request_load(client, total, concurrency, timeout, stability)
                elapsed = (time.perf_counter() - started) * 1000.0
                summarize("roundtrip", f"concurrency={concurrency}", transport, latencies, elapsed, 1, stability, mem_mb())
            finally:
                with suppress(Exception):
                    await asyncio.wait_for(client.close(), 2.0)
                if client.communicator.is_alive():
                    stability.pending_leak += 1
        except Exception as exc:  # noqa: BLE001
            stability.timeouts += 1
            log(f"[roundtrip] concurrency={concurrency} error={exc}")
            summarize("roundtrip", f"concurrency={concurrency}", transport, [], 0.0, 1, stability, mem_mb())
        finally:
            await server.stop()
        log(f"[roundtrip] transport={transport} concurrency={concurrency} done violations={stability.violations()}")


async def profile_burst(transport: str, mode: str) -> None:
    counts = [200] if mode == "quick" else [1000, 10000, 100000]
    log(f"[burst] transport={transport} mode={mode} counts={counts}")
    for count in counts:
        stability = Stability()
        state = {"received": 0, "last_index": -1}
        loop = asyncio.get_running_loop()
        done = loop.create_future()

        def on_connected(client):
            def on_message(data):
                if data.index <= state["last_index"]:
                    stability.fifo_violations += 1
                state["last_index"] = data.index
                state["received"] += 1
                if state["received"] >= count and not done.done():
                    done.set_result(None)

            client.communicator.add_listener(type_name_of(TestData), on_message)

        port = free_port()
        server = make_server(transport, port, on_connected)
        await server.start()
        client = None
        try:
            client = await dial_with_retry(transport, port)
            started = time.perf_counter()
            for i in range(count):
                await client.communicator.send(TestData(index=i, message=f"b-{i}"))
            try:
                await asyncio.wait_for(done, 10.0 if mode == "quick" else 60.0)
            except TimeoutError:
                stability.timeouts += 1
                log(f"[burst] count={count} dispatch timeout received={state['received']}")
            elapsed = (time.perf_counter() - started) * 1000.0
            if state["received"] != count:
                stability.pending_leak += 1
            emit_throughput("burst", f"count={count}", transport, count, elapsed, 1, test_data_payload_bytes("b-0"), stability, mem_mb())
        except Exception as exc:  # noqa: BLE001
            stability.timeouts += 1
            log(f"[burst] count={count} error={exc}")
            emit_throughput("burst", f"count={count}", transport, state["received"], 0.0, 1, test_data_payload_bytes("b-0"), stability, mem_mb())
        finally:
            if client is not None:
                with suppress(Exception):
                    await asyncio.wait_for(client.close(), 2.0)
                if client.communicator.is_alive():
                    stability.pending_leak += 1
            await server.stop()
        log(f"[burst] transport={transport} count={count} received={state['received']} violations={stability.violations()}")


async def profile_sustained(transport: str, mode: str) -> None:
    durations_ms = [2000] if mode == "quick" else [30000, 300000, 1800000]
    concurrency = 16
    timeout = 15.0
    log(f"[sustained] transport={transport} mode={mode} durations={durations_ms}ms concurrency={concurrency}")
    for duration_ms in durations_ms:
        stability = Stability()
        peak = mem_mb()
        port = free_port()
        server = make_echo_server(transport, port)
        await server.start()
        try:
            client = await dial_with_retry(transport, port)
            latencies: list[float] = []
            next_index = 0
            loop = asyncio.get_running_loop()
            deadline = loop.time() + duration_ms / 1000.0
            try:
                while loop.time() < deadline:
                    indices = list(range(next_index, next_index + concurrency))
                    next_index += concurrency

                    async def one(index: int) -> None:
                        started = time.perf_counter()
                        try:
                            res = await client.communicator.send_request(
                                TestData(index=index, message=f"s-{index}"), TestData, timeout=timeout
                            )
                            latencies.append((time.perf_counter() - started) * 1000.0)
                            if res.index != index * 2 or res.message != f"echo:s-{index}":
                                stability.nonce_mismatch += 1
                        except Exception as exc:  # noqa: BLE001
                            classify_request_error(str(exc), stability)

                    await asyncio.gather(*(one(i) for i in indices))
                    current = mem_mb()
                    if current > peak:
                        peak = current
                summarize("sustained", f"duration={duration_ms}ms", transport, latencies, float(duration_ms), 1, stability, peak)
            finally:
                with suppress(Exception):
                    await asyncio.wait_for(client.close(), 2.0)
                if client.communicator.is_alive():
                    stability.pending_leak += 1
        except Exception as exc:  # noqa: BLE001
            stability.timeouts += 1
            log(f"[sustained] duration={duration_ms}ms error={exc}")
            summarize("sustained", f"duration={duration_ms}ms", transport, [], 0.0, 1, stability, peak)
        finally:
            await server.stop()
        log(f"[sustained] transport={transport} duration={duration_ms}ms done peakMemMb={peak:.1f} violations={stability.violations()}")


async def profile_parallel(transport: str, mode: str) -> None:
    client_counts = [4] if mode == "quick" else [16, 128, 512, 1024]
    per_client = 50 if mode == "quick" else 500
    concurrency = 8
    timeout = 15.0
    log(f"[parallel] transport={transport} mode={mode} clients={client_counts} perClient={per_client}")
    for client_count in client_counts:
        stability = Stability()
        port = free_port()
        server = make_echo_server(transport, port)
        await server.start()
        clients: list = []
        try:
            for _ in range(client_count):
                try:
                    clients.append(await dial_with_retry(transport, port))
                except Exception as exc:  # noqa: BLE001
                    stability.timeouts += 1
                    log(f"[parallel] clients={client_count} dial error={exc}")
            all_latencies: list[float] = []
            started = time.perf_counter()

            async def run_one(client) -> None:
                all_latencies.extend(await run_request_load(client, per_client, concurrency, timeout, stability))

            await asyncio.gather(*(run_one(c) for c in clients))
            elapsed = (time.perf_counter() - started) * 1000.0
            summarize("parallel", f"clients={client_count}", transport, all_latencies, elapsed, client_count, stability, mem_mb())
        finally:
            for client in clients:
                with suppress(Exception):
                    await asyncio.wait_for(client.close(), 2.0)
                if client.communicator.is_alive():
                    stability.pending_leak += 1
            await server.stop()
        log(f"[parallel] transport={transport} clients={client_count} done violations={stability.violations()}")


async def profile_slow_mix(transport: str, mode: str) -> None:
    client_count = 4 if mode == "quick" else 16
    per_client = 20 if mode == "quick" else 200
    slow_every = 5 if mode == "quick" else 10
    delay_ms = 50 if mode == "quick" else 100
    timeout = 10.0 if mode == "quick" else 30.0
    axis = f"clients={client_count},slowEvery={slow_every},delayMs={delay_ms}"
    stability = Stability()
    peak = mem_mb()
    latencies: list[float] = []
    elapsed = 0.0
    log(f"[slow-mix] transport={transport} mode={mode} {axis} perClient={per_client}")

    port = free_port()
    server = make_slow_mix_server(transport, port, slow_every, delay_ms / 1000.0)
    await server.start()
    clients: list[tuple[int, object]] = []
    try:
        for client_id in range(client_count):
            try:
                clients.append((client_id, await dial_with_retry(transport, port)))
            except Exception as exc:  # noqa: BLE001
                stability.timeouts += 1
                log(f"[slow-mix] client={client_id} dial error={exc}")

        completed_by_client = [-1] * client_count
        started = time.perf_counter()

        async def one_request(client, client_id: int, index: int) -> None:
            message = f"sm-{client_id}-{index}"
            req_start = time.perf_counter()
            try:
                res = await client.communicator.send_request(
                    TestData(index=index, message=message), TestData, timeout=timeout
                )
                latencies.append((time.perf_counter() - req_start) * 1000.0)
                if res.index != index * 2 or res.message != "echo:" + message:
                    stability.nonce_mismatch += 1
                if index <= completed_by_client[client_id]:
                    stability.fifo_violations += 1
                completed_by_client[client_id] = index
            except Exception as exc:  # noqa: BLE001
                classify_request_error(str(exc), stability)

        await asyncio.gather(
            *(
                one_request(client, client_id, index)
                for client_id, client in clients
                for index in range(per_client)
            )
        )
        peak = max(peak, mem_mb())
        elapsed = (time.perf_counter() - started) * 1000.0
    finally:
        for _, client in clients:
            with suppress(Exception):
                await asyncio.wait_for(client.close(), 2.0)
            if client.communicator.is_alive():
                stability.pending_leak += 1
        await server.stop()
    ordered = sorted(latencies)
    throughput = (len(latencies) / elapsed * 1000.0) if elapsed > 0 else 0.0
    emit_row(
        profile="slow-mix",
        axis=axis,
        transport=transport,
        payload_bytes=test_data_payload_bytes("sm-0-0"),
        client_count=client_count,
        requests=len(latencies),
        throughput=throughput,
        p50=percentile(ordered, 50),
        p95=percentile(ordered, 95),
        p99=percentile(ordered, 99),
        mem=peak,
        stability=stability,
    )
    log(
        f"[slow-mix] transport={transport} clients={client_count} done "
        f"peakMemMb={peak:.1f} violations={stability.violations()}"
    )


async def profile_payload(transport: str, mode: str) -> None:
    sizes = PAYLOAD_SIZES_QUICK if mode == "quick" else PAYLOAD_SIZES_FULL
    concurrency = 4 if mode == "quick" else 8
    requests = 20 if mode == "quick" else 100
    timeout = 10.0 if mode == "quick" else 30.0
    log(
        f"[payload] transport={transport} mode={mode} sizes={[payload_label(s) for s in sizes]} "
        f"concurrency={concurrency} requests/size={requests}"
    )
    for size in sizes:
        stability = Stability()
        filler = payload_filler(size)
        expected_echo = "echo:" + filler
        payload_bytes = test_data_payload_bytes(filler)
        axis = f"payload={payload_label(size)}"
        peak = mem_mb()
        port = free_port()
        server = make_echo_server(transport, port)
        await server.start()
        try:
            client = await dial_with_retry(transport, port)
            try:
                latencies: list[float] = []
                issued = 0
                next_index = 0
                started = time.perf_counter()

                async def one(index: int) -> None:
                    req_start = time.perf_counter()
                    try:
                        res = await client.communicator.send_request(
                            TestData(index=index, message=filler), TestData, timeout=timeout
                        )
                        latencies.append((time.perf_counter() - req_start) * 1000.0)
                        if res.index != index * 2 or res.message != expected_echo:
                            stability.nonce_mismatch += 1
                    except Exception as exc:  # noqa: BLE001
                        classify_request_error(str(exc), stability)

                while issued < requests:
                    batch_size = min(concurrency, requests - issued)
                    indices = list(range(next_index, next_index + batch_size))
                    next_index += batch_size
                    issued += batch_size
                    await asyncio.gather(*(one(i) for i in indices))
                    current = mem_mb()
                    if current > peak:
                        peak = current
                elapsed = (time.perf_counter() - started) * 1000.0
                ordered = sorted(latencies)
                throughput = (len(latencies) / elapsed * 1000.0) if elapsed > 0 else 0.0
                emit_row(
                    profile="payload",
                    axis=axis,
                    transport=transport,
                    payload_bytes=payload_bytes,
                    client_count=1,
                    requests=len(latencies),
                    throughput=throughput,
                    p50=percentile(ordered, 50),
                    p95=percentile(ordered, 95),
                    p99=percentile(ordered, 99),
                    mem=peak,
                    stability=stability,
                )
            finally:
                with suppress(Exception):
                    await asyncio.wait_for(client.close(), 2.0)
                if client.communicator.is_alive():
                    stability.pending_leak += 1
        except Exception as exc:  # noqa: BLE001
            stability.timeouts += 1
            log(f"[payload] size={payload_label(size)} error={exc}")
            emit_row(
                profile="payload",
                axis=axis,
                transport=transport,
                payload_bytes=payload_bytes,
                client_count=1,
                requests=0,
                throughput=0.0,
                p50=0.0,
                p95=0.0,
                p99=0.0,
                mem=peak,
                stability=stability,
            )
        finally:
            await server.stop()
        log(
            f"[payload] transport={transport} size={payload_label(size)} bytes={payload_bytes} "
            f"peakMemMb={peak:.1f} violations={stability.violations()}"
        )


RUNNERS = {
    "roundtrip": profile_roundtrip,
    "burst": profile_burst,
    "sustained": profile_sustained,
    "parallel": profile_parallel,
    "slow-mix": profile_slow_mix,
    "payload": profile_payload,
}


def parse_list(value: str | None, default: list[str]) -> list[str]:
    if not value:
        return default
    items = [v.strip() for v in value.split(",") if v.strip()]
    return items or default


async def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["quick", "full"], default="quick")
    parser.add_argument("--quick", action="store_const", const="quick", dest="mode")
    parser.add_argument("--full", action="store_const", const="full", dest="mode")
    parser.add_argument("--transport", "--transports", dest="transport", default=None)
    parser.add_argument("--profiles", "--profile", dest="profiles", default=None)
    args = parser.parse_args()

    mode = args.mode or "quick"
    transports = parse_list(args.transport, ["tcp", "ws"])
    profiles = parse_list(args.profiles, ALL_PROFILES)

    log(
        f"INFO stress harness language={LANGUAGE} mode={mode} transports={','.join(transports)} "
        f"profiles={','.join(profiles)} typeName={type_name_of(TestData)}"
    )

    for transport in transports:
        if transport not in ("tcp", "ws"):
            emit_skip("all", "transport", transport, "unsupported transport for Python same-language baseline")
            continue
        for profile in profiles:
            runner = RUNNERS.get(profile)
            if runner is None:
                emit_skip(profile, "profile", transport, "profile not part of Python same-language baseline (gateway is TypeScript-specific)")
                continue
            await runner(transport, mode)

    status = "PASS" if _total_violations == 0 else "FAIL"
    print(
        f"SUMMARY|status={status}|language={LANGUAGE}|mode={mode}|transports={','.join(transports)}|"
        f"profiles={','.join(profiles)}|rows={_rows_emitted}|stability_violations={_total_violations}",
        flush=True,
    )
    if status != "PASS":
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())