#!/usr/bin/env python3
"""
pipeview-tauri 极限压力测试脚本

模拟各种高数据量场景，测试 Tauri 前端的渲染性能和稳定性：
  - 文本洪水：大文本行 + ANSI 颜色
  - 十六进制洪水：大批量 hex dump
  - 波形洪水：高频 plot 采样点
  - 混合模式：同时发送文本 + 波形（MixedTextPlot）
  - 突发模式：间歇性峰值流量

用法:
  python tools/stress_test.py --mode text --rate 5000 --port 8091
  python tools/stress_test.py --mode plot --rate 200 --channels 8 --port 8092
  python tools/stress_test.py --mode mixed --rate 1000 --port 8093
  python tools/stress_test.py --mode burst --rate 10000 --burst-size 500 --port 8094
"""

import argparse
import math
import random
import socket
import struct
import sys
import threading
import time
from collections import defaultdict

PLOT_ESCAPE = 0x1E
PLOT_MARKER = ord("P")

DISCONNECT_WINERRORS = {
    10053,  # Software caused connection abort
    10054,  # Connection reset by peer
    10058,  # Socket shutdown race on Windows
}

# ── ANSI Color Palette ─────────────────────────────────────────────

ANSI_COLORS = [
    "\033[31m", "\033[32m", "\033[33m", "\033[34m", "\033[35m", "\033[36m",
    "\033[91m", "\033[92m", "\033[93m", "\033[94m", "\033[95m", "\033[96m",
    "\033[1;31m", "\033[1;32m", "\033[1;33m", "\033[1;34m",
]
ANSI_RESET = "\033[0m"

# ── Lorem ipsum lines for text stress testing ──────────────────────

LOREM_LINES = [
    "[INFO] Sensor data received: temperature={temp:.2f}°C humidity={hum:.1f}% pressure={pres:.1f}hPa",
    "[DEBUG] Frame #{n:06d} decoded pipeline={pipe} latency={lat:.3f}ms",
    "[WARN] Buffer utilization {pct:.1f}% — threshold approaching",
    "[ERROR] Checksum mismatch at offset {off:#x} expected={exp:#x} got={got:#x}",
    "[TRACE] I2C transaction: addr={addr:#x} reg={reg:#x} val={val:#x}",
    "[METRIC] throughput={tput:.1f} lines/s memory={mem:.1f}MB active_sessions={sess}",
    "[EVENT] Session {sid} state changed: {old} -> {new}",
    "[DATA] {ts} | CH{ch:02d} | {v0:+08.4f} | {v1:+08.4f} | {v2:+08.4f}",
]

# ── Text stress generator ───────────────────────────────────────────

def generate_text_line(seq: int, ansi: bool = False, payload_size: int = 80) -> bytes:
    """Generate a single text line with optional ANSI colors."""
    template = random.choice(LOREM_LINES)
    line = template.format(
        temp=20.0 + 10 * math.sin(seq * 0.1),
        hum=45.0 + 20 * math.cos(seq * 0.07),
        pres=1013.0 + random.uniform(-5, 5),
        n=seq, pipe=f"pipe_{seq % 4}", lat=random.uniform(0.1, 5.0),
        pct=random.uniform(10, 95), off=seq * 16,
        exp=random.randint(0, 255), got=random.randint(0, 255),
        addr=0x40 + (seq % 8), reg=0x00 + (seq % 16), val=random.randint(0, 65535),
        tput=random.uniform(100, 10000), mem=random.uniform(50, 500), sess=random.randint(1, 8),
        sid=seq % 8 + 1, old="disconnected", new="connected",
        ts=time.strftime("%H:%M:%S", time.localtime()),
        ch=seq % 8, v0=random.uniform(-10, 10), v1=random.uniform(-10, 10), v2=random.uniform(-10, 10),
    )

    # Pad or trim to target payload size
    if len(line) < payload_size:
        line += " " + "x" * (payload_size - len(line) - 1)
    elif len(line) > payload_size:
        line = line[:payload_size]

    if ansi:
        color = random.choice(ANSI_COLORS)
        line = f"{color}{line}{ANSI_RESET}"

    return (line + "\n").encode("utf-8")


def generate_hex_line(seq: int, bytes_per_line: int = 32) -> bytes:
    """Generate a hex dump line (as text, mimicking hex view data)."""
    data = bytes([(seq * bytes_per_line + i) % 256 for i in range(bytes_per_line)])
    hex_part = " ".join(f"{b:02X}" for b in data)
    ascii_part = "".join(chr(b) if 32 <= b < 127 else "." for b in data)
    return f"[{seq:06d}] {hex_part}  |{ascii_part}|\n".encode("utf-8")


# ── Plot stress generator ───────────────────────────────────────────

def build_plot_frame(
    sample_index: int,
    channels: int,
    plot_format: str,
    samples_per_channel: int,
    amplitude: float,
    frequency_hz: float,
    sample_rate_hz: float,
) -> bytes:
    """Build a raw COBS-encoded plot frame (matching test_plot.py format)."""
    values = []
    if plot_format == "xy":
        for _ in range(samples_per_channel):
            t = (sample_index * samples_per_channel + len(values)) / sample_rate_hz
            values.append(amplitude * math.sin(2 * math.pi * frequency_hz * t))
            values.append(amplitude * math.cos(2 * math.pi * frequency_hz * t))
    elif plot_format == "block":
        for ch in range(channels):
            phase = 2 * math.pi * ch / channels
            for _ in range(samples_per_channel):
                t = (sample_index * samples_per_channel + len(values) // channels) / sample_rate_hz
                values.append(amplitude * math.sin(2 * math.pi * frequency_hz * t + phase))
    else:  # interleaved
        for _ in range(samples_per_channel):
            t = (sample_index * samples_per_channel + len(values) // channels) / sample_rate_hz
            for ch in range(channels):
                phase = 2 * math.pi * ch / channels
                values.append(amplitude * math.sin(2 * math.pi * frequency_hz * t + phase))

    # Pack as f32 little-endian
    payload = struct.pack(f"<{len(values)}f", *values)
    return cobs_encode(payload)


def cobs_encode(data: bytes) -> bytes:
    """Consistent Overhead Byte Stuffing encoder."""
    result = bytearray()
    block_start = 0
    while block_start < len(data):
        end = min(block_start + 254, len(data))
        block = data[block_start:end]
        if end < len(data):
            result.append(len(block) + 1)
            result.extend(block)
        else:
            result.append(len(block) + 1 if len(block) < 254 else 255)
            result.extend(block)
            result.append(0)
        block_start = end
    return bytes(result)


def build_mixed_frame(seq: int, text_rate_per_plot: int, channels: int) -> bytes:
    """Build a MixedTextPlot frame: text lines + one COBS plot frame."""
    frames = []
    for i in range(text_rate_per_plot):
        frames.append(generate_text_line(seq * text_rate_per_plot + i, ansi=True))
    # Add one plot frame per text_rate_per_plot text lines
    plot_data = build_plot_frame(seq, channels, "interleaved", 32, 100.0, 10.0, 1000.0)
    frames.append(bytes([PLOT_ESCAPE, PLOT_MARKER]) + plot_data)
    return b"".join(frames)


# ── TCP Server ──────────────────────────────────────────────────────

class StressServer:
    def __init__(self, host: str, port: int, mode: str, rate: float,
                 payload_size: int, channels: int, burst_size: int,
                 duration: float, ansi: bool):
        self.host = host
        self.port = port
        self.mode = mode
        self.rate = rate          # lines or frames per second
        self.payload_size = payload_size
        self.channels = channels
        self.burst_size = burst_size
        self.duration = duration
        self.ansi = ansi
        self.stats = defaultdict(int)
        self.running = False
        self.clients = []

    def start(self):
        self.running = True
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind((self.host, self.port))
        sock.listen(5)
        sock.settimeout(1.0)

        print(f"[STRESS] {self.mode.upper()} mode | {self.rate} lines/s")
        print(f"[STRESS] Listening on {self.host}:{self.port}")
        if self.duration > 0:
            print(f"[STRESS] Duration: {self.duration}s")
        print(f"[STRESS] Payload: {self.payload_size}B | Channels: {self.channels}")
        if self.mode == "burst":
            print(f"[STRESS] Burst size: {self.burst_size} lines")
        print()

        stats_thread = threading.Thread(target=self._print_stats, daemon=True)
        stats_thread.start()

        accept_thread = threading.Thread(target=self._accept_loop, args=(sock,), daemon=True)
        accept_thread.start()

        try:
            while self.running:
                time.sleep(0.1)
        except KeyboardInterrupt:
            print("\n[STRESS] Shutting down...")
        finally:
            self.running = False
            sock.close()

    def _accept_loop(self, sock):
        while self.running:
            try:
                conn, addr = sock.accept()
                print(f"[STRESS] Client connected: {addr[0]}:{addr[1]}")
                t = threading.Thread(target=self._client_loop, args=(conn, addr), daemon=True)
                t.start()
                self.clients.append(t)
            except socket.timeout:
                continue
            except OSError:
                break

    def _client_loop(self, conn: socket.socket, addr):
        seq = 0
        start_time = time.time()
        last_count_time = start_time
        local_count = 0

        # Pre-compute interval
        if self.mode == "burst":
            interval = 0
            burst_interval = max(0.016, self.burst_size / max(self.rate, 1))
        else:
            interval = 1.0 / max(self.rate, 1) if self.rate > 0 else 0.016

        try:
            conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            while self.running:
                elapsed = time.time() - start_time
                if self.duration > 0 and elapsed >= self.duration:
                    break

                if self.mode == "burst":
                    self._send_burst(conn, seq)
                    seq += self.burst_size
                    local_count += self.burst_size
                    time.sleep(burst_interval)
                else:
                    data = self._generate(seq)
                    conn.sendall(data)
                    seq += 1
                    local_count += 1

                    if interval > 0:
                        # Spin-wait for precise timing at high rates
                        target = start_time + (seq / self.rate)
                        sleep_time = target - time.time()
                        if sleep_time > 0:
                            time.sleep(min(sleep_time, interval))

                # Periodic stats update
                if time.time() - last_count_time >= 1.0:
                    with threading.Lock():
                        self.stats["lines"] += local_count
                        self.stats["bytes"] += local_count * self.payload_size  # approximate
                    local_count = 0
                    last_count_time = time.time()

        except OSError as e:
            if hasattr(e, "winerror") and e.winerror in DISCONNECT_WINERRORS:
                pass
            elif not self.running:
                pass
            else:
                self.stats["errors"] += 1
        finally:
            try:
                conn.close()
            except OSError:
                pass
            print(f"[STRESS] Client disconnected: {addr[0]}:{addr[1]}")

    def _send_burst(self, conn, start_seq):
        """Send burst_size lines as fast as possible."""
        data = b"".join(self._generate(start_seq + i) for i in range(self.burst_size))
        conn.sendall(data)

    def _generate(self, seq: int) -> bytes:
        if self.mode == "text":
            return generate_text_line(seq, ansi=self.ansi, payload_size=self.payload_size)
        elif self.mode == "hex":
            return generate_hex_line(seq, bytes_per_line=max(4, self.payload_size // 3))
        elif self.mode == "plot":
            return build_plot_frame(
                seq, self.channels, "interleaved",
                samples_per_channel=32, amplitude=100.0,
                frequency_hz=10.0, sample_rate_hz=1000.0,
            )
        elif self.mode == "mixed":
            return build_mixed_frame(seq, text_rate_per_plot=5, channels=self.channels)
        else:
            return generate_text_line(seq, ansi=True, payload_size=self.payload_size)

    def _print_stats(self):
        while self.running:
            time.sleep(2.0)
            with threading.Lock():
                lines = self.stats["lines"]
                errors = self.stats["errors"]
                b = self.stats["bytes"]
            if lines > 0:
                kb_s = b / 2048  # KB/s over 2 seconds
                print(f"[STATS] {lines:>8d} lines | {kb_s:>8.1f} KB/s | {errors:>4d} errors")
            self.stats.clear()


# ── CLI ─────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="pipeview-tauri 极限压力测试",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 文本洪水：5000 行/秒，1KB 行宽，ANSI 彩色
  python tools/stress_test.py --mode text --rate 5000 --payload 1024 --ansi

  # 波形洪水：200 帧/秒，8 通道
  python tools/stress_test.py --mode plot --rate 200 --channels 8 --port 8092

  # 突发模式：每秒爆发一次 2000 行峰值
  python tools/stress_test.py --mode burst --rate 10000 --burst-size 2000

  # 混合模式：文本 + 波形（MixedTextPlot 帧）
  python tools/stress_test.py --mode mixed --rate 1000 --channels 4

  # 压测 60 秒后自动停止
  python tools/stress_test.py --mode text --rate 10000 --duration 60
        """,
    )
    parser.add_argument("--host", default="127.0.0.1", help="绑定地址 (default: 127.0.0.1)")
    parser.add_argument("--port", type=int, default=8091, help="端口 (default: 8091)")
    parser.add_argument(
        "--mode", choices=["text", "hex", "plot", "mixed", "burst"],
        default="text", help="测试模式 (default: text)"
    )
    parser.add_argument(
        "--rate", type=float, default=1000,
        help="目标速率 (lines/s 或 frames/s, default: 1000)"
    )
    parser.add_argument(
        "--payload", type=int, default=256,
        help="文本行/hex 行目标字节数 (default: 256)"
    )
    parser.add_argument(
        "--channels", type=int, default=4,
        help="Plot 通道数 (default: 4)"
    )
    parser.add_argument(
        "--burst-size", type=int, default=500,
        help="突发模式每次发送的行数 (default: 500)"
    )
    parser.add_argument(
        "--duration", type=float, default=0,
        help="运行时长（秒），0 表示无限制 (default: 0)"
    )
    parser.add_argument(
        "--ansi", action="store_true",
        help="文本模式启用 ANSI 颜色"
    )

    args = parser.parse_args()

    if args.mode == "plot" and args.rate > 500:
        print("[WARN] Plot rate > 500 may overwhelm the renderer. Consider --rate 100-200.")
        print()

    server = StressServer(
        host=args.host,
        port=args.port,
        mode=args.mode,
        rate=args.rate,
        payload_size=args.payload,
        channels=args.channels,
        burst_size=args.burst_size,
        duration=args.duration,
        ansi=args.ansi,
    )
    server.start()


if __name__ == "__main__":
    main()