"""
パフォーマンスベンチマークスクリプト

ZeroMQ PUB/SUBの通信スループットとレイテンシを計測する。
- スループット: 最大 msg/s
- レイテンシ: avg / p50 / p95 / p99 / max
- メッセージロス率
"""

import sys
import time
import threading
import statistics
from collections import deque

from rich.console import Console
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn

from rithmic_bridge.config import BridgeConfig, ZmqConfig
from rithmic_bridge.models import (
    MBOEvent, TradeEvent, HeartbeatEvent,
    MBOAction, Side, TradeAction,
    TOPIC_MBO, TOPIC_TRADE, TOPIC_HEARTBEAT,
)
from rithmic_bridge.publisher.server import ZmqPublisher
from rithmic_bridge.client.receiver import ZmqReceiver

console = Console()


def run_throughput_benchmark(
    target_rate: int = 5000,
    duration_seconds: int = 10,
    port: int = 15700,
) -> dict:
    """
    スループットベンチマーク

    指定レートでMBOイベントを送信し、受信側のスループットとロス率を計測する。

    Args:
        target_rate: 目標送信レート (msg/s)
        duration_seconds: テスト時間 (秒)
        port: 使用するポート
    """
    config = BridgeConfig(
        zmq=ZmqConfig(
            pub_addr=f"tcp://127.0.0.1:{port}",
            pub_port=port,
            client_host="127.0.0.1",
            client_port=port,
            hwm=50000,  # ベンチマーク用に大きめのバッファ
        )
    )

    # 受信カウンター
    recv_count = [0]
    recv_lock = threading.Lock()

    def on_mbo(event: MBOEvent):
        with recv_lock:
            recv_count[0] += 1

    # パブリッシャー & レシーバー起動
    publisher = ZmqPublisher(config)
    publisher.start(enable_heartbeat=False)
    time.sleep(0.2)

    receiver = ZmqReceiver(config)
    receiver.on_mbo(on_mbo)
    receiver.start(topics=[TOPIC_MBO])
    time.sleep(0.5)

    # 送信フェーズ
    total_sent = 0
    interval = 1.0 / target_rate if target_rate > 0 else 0
    # バースト送信用のバッチサイズ
    batch_size = max(1, target_rate // 100)

    start_time = time.time()
    end_time = start_time + duration_seconds

    with Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]スループット計測中..."),
        BarColumn(),
        TextColumn("{task.percentage:>3.0f}%"),
        TextColumn("送信: {task.fields[sent]} msg"),
        console=console,
    ) as progress:
        task = progress.add_task(
            "benchmark",
            total=duration_seconds * 10,  # 0.1秒単位で更新
            sent=0,
        )

        while time.time() < end_time:
            # バッチ送信
            for _ in range(batch_size):
                event = MBOEvent(
                    symbol="GCJ6.COMEX",
                    action=MBOAction.ADD,
                    order_id=str(total_sent),
                    side=Side.BID,
                    price=2500.0,
                    size=5,
                )
                publisher.enqueue_mbo(event.to_json())
                total_sent += 1

            # レート制御
            if interval > 0:
                time.sleep(interval * batch_size)

            # プログレスバー更新
            elapsed = time.time() - start_time
            progress.update(
                task,
                completed=int(elapsed * 10),
                sent=total_sent,
            )

    # 受信完了待ち
    time.sleep(2.0)

    actual_elapsed = time.time() - start_time
    with recv_lock:
        total_received = recv_count[0]

    # クリーンアップ
    receiver.stop()
    publisher.stop()

    # 結果集計
    send_rate = total_sent / actual_elapsed if actual_elapsed > 0 else 0
    recv_rate = total_received / actual_elapsed if actual_elapsed > 0 else 0
    loss_rate = 1 - (total_received / total_sent) if total_sent > 0 else 0

    return {
        "target_rate": target_rate,
        "total_sent": total_sent,
        "total_received": total_received,
        "duration": actual_elapsed,
        "send_rate": send_rate,
        "recv_rate": recv_rate,
        "loss_rate": loss_rate,
    }


def run_latency_benchmark(
    num_messages: int = 1000,
    port: int = 15701,
) -> dict:
    """
    レイテンシベンチマーク

    MBOイベントをゆっくり送信し、各メッセージのレイテンシを計測する。
    """
    config = BridgeConfig(
        zmq=ZmqConfig(
            pub_addr=f"tcp://127.0.0.1:{port}",
            pub_port=port,
            client_host="127.0.0.1",
            client_port=port,
            hwm=10000,
        )
    )

    # レイテンシ計測リスト
    latencies = []  # ミリ秒
    recv_lock = threading.Lock()

    def on_mbo(event: MBOEvent):
        recv_time = time.time()
        latency_ms = (recv_time - event.timestamp) * 1000.0
        with recv_lock:
            latencies.append(latency_ms)

    # パブリッシャー & レシーバー起動
    publisher = ZmqPublisher(config)
    publisher.start(enable_heartbeat=False)
    time.sleep(0.2)

    receiver = ZmqReceiver(config)
    receiver.on_mbo(on_mbo)
    receiver.start(topics=[TOPIC_MBO])
    time.sleep(0.5)

    # 送信フェーズ（100 msg/s で均等送信）
    send_rate = 100  # msg/s
    interval = 1.0 / send_rate

    with Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]レイテンシ計測中..."),
        BarColumn(),
        TextColumn("{task.percentage:>3.0f}%"),
        console=console,
    ) as progress:
        task = progress.add_task(
            "latency",
            total=num_messages,
        )

        for i in range(num_messages):
            event = MBOEvent(
                symbol="GCJ6.COMEX",
                action=MBOAction.ADD,
                order_id=str(i),
                side=Side.BID,
                price=2500.0,
                size=5,
                # タイムスタンプは送信直前に設定
                timestamp=time.time(),
            )
            publisher.enqueue_mbo(event.to_json())
            time.sleep(interval)
            progress.update(task, completed=i + 1)

    # 受信完了待ち
    time.sleep(1.0)

    # クリーンアップ
    receiver.stop()
    publisher.stop()

    # 結果集計（負値はキュー遅延によるアーティファクトなので除外）
    with recv_lock:
        received_latencies = [lat for lat in latencies if lat >= 0]

    if not received_latencies:
        return {
            "num_sent": num_messages,
            "num_received": 0,
            "avg": 0, "p50": 0, "p95": 0, "p99": 0, "max": 0, "min": 0,
        }

    received_latencies.sort()
    n = len(received_latencies)

    return {
        "num_sent": num_messages,
        "num_received": n,
        "avg": statistics.mean(received_latencies),
        "p50": received_latencies[n // 2],
        "p95": received_latencies[int(n * 0.95)],
        "p99": received_latencies[int(n * 0.99)],
        "max": max(received_latencies),
        "min": min(received_latencies),
    }


def main():
    """ベンチマークメインエントリーポイント"""
    console.print("\n[bold blue]═══════════════════════════════════════════════[/]")
    console.print("[bold blue]  Rithmic Bridge パフォーマンスベンチマーク[/]")
    console.print("[bold blue]═══════════════════════════════════════════════[/]\n")

    # === 1. スループットベンチマーク ===
    console.print("[bold cyan]1. スループットベンチマーク[/]\n")

    # 複数レートでテスト
    throughput_results = []
    for rate in [1000, 5000, 10000, 20000]:
        console.print(f"  [dim]目標: {rate:,} msg/s ...[/]")
        result = run_throughput_benchmark(
            target_rate=rate,
            duration_seconds=5,
            port=15700 + len(throughput_results),
        )
        throughput_results.append(result)
        time.sleep(0.5)  # ポート解放待ち

    # 結果テーブル
    table = Table(title="スループット結果")
    table.add_column("目標 (msg/s)", style="cyan", justify="right")
    table.add_column("送信 (msg/s)", style="white", justify="right")
    table.add_column("受信 (msg/s)", style="green", justify="right")
    table.add_column("送信数", style="white", justify="right")
    table.add_column("受信数", style="white", justify="right")
    table.add_column("ロス率", style="yellow", justify="right")

    for r in throughput_results:
        loss_color = "green" if r["loss_rate"] < 0.01 else "yellow" if r["loss_rate"] < 0.1 else "red"
        table.add_row(
            f"{r['target_rate']:,}",
            f"{r['send_rate']:,.0f}",
            f"{r['recv_rate']:,.0f}",
            f"{r['total_sent']:,}",
            f"{r['total_received']:,}",
            f"[{loss_color}]{r['loss_rate']:.1%}[/]",
        )

    console.print(table)
    console.print()

    # === 2. レイテンシベンチマーク ===
    console.print("[bold cyan]2. レイテンシベンチマーク[/]\n")
    lat_result = run_latency_benchmark(num_messages=1000, port=15710)

    lat_table = Table(title="レイテンシ結果 (ms)")
    lat_table.add_column("指標", style="cyan")
    lat_table.add_column("値", style="green", justify="right")

    lat_table.add_row("送信数", str(lat_result["num_sent"]))
    lat_table.add_row("受信数", str(lat_result["num_received"]))

    avg_color = "green" if lat_result["avg"] < 1 else "yellow" if lat_result["avg"] < 5 else "red"
    lat_table.add_row("平均", f"[{avg_color}]{lat_result['avg']:.3f}[/]")
    lat_table.add_row("P50", f"{lat_result['p50']:.3f}")
    lat_table.add_row("P95", f"{lat_result['p95']:.3f}")
    lat_table.add_row("P99", f"{lat_result['p99']:.3f}")
    lat_table.add_row("最小", f"{lat_result['min']:.3f}")
    lat_table.add_row("最大", f"{lat_result['max']:.3f}")

    console.print(lat_table)

    # === 総合結果 ===
    console.print("\n[bold blue]═══════════════════════════════════════════════[/]")
    console.print("[bold blue]  ベンチマーク完了[/]")
    console.print("[bold blue]═══════════════════════════════════════════════[/]\n")


if __name__ == "__main__":
    main()
