"""Local-only benchmark: drive the file-helper binary directly over a pipe.

No SSH, no Docker. Establishes a protocol-overhead baseline before Phase B/C
add real network/container latency.

Beyond the minimal single-read timings from the design plan, this script also
measures sustained sequential throughput, stat/list latency, and concurrent
request behaviour, so Phase A can be checked against the plan's performance
gate (>=20 MiB/s, or clearly "not a future bottleneck") before any SSH/Docker
work begins. Run this script multiple times and compare/median the numbers;
a single run is not a statistically meaningful benchmark.
"""

from __future__ import annotations

import os
import shutil
import subprocess
import sys
import tempfile
import threading
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from msdev.core.request_client import RequestClient
from msdev.core import streamproto as proto

HELPER_BINARY = os.path.join(
    os.path.dirname(__file__), "..", "file-helper", "bin", "msdev-file-helper-dev",
)

# Large enough to make one-shot process/pipe setup overhead negligible while
# still finishing in well under a second at the throughput this protocol
# should sustain locally.
THROUGHPUT_FILE_BYTES = 64 * 1024 * 1024
THROUGHPUT_CHUNK_BYTES = 4 * 1024 * 1024
STAT_CALL_COUNT = 200
LIST_DIRECTORY_ENTRY_COUNT = 64
CONCURRENT_READERS = 8
CONCURRENT_READ_BYTES_PER_READER = 8 * 1024 * 1024
CONCURRENT_CHUNK_BYTES = 1 * 1024 * 1024


def _start_helper() -> tuple[subprocess.Popen, RequestClient]:
    process = subprocess.Popen(
        [HELPER_BINARY], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
    )

    def abort_transport() -> None:
        if process.poll() is None:
            process.terminate()

    client = RequestClient(
        process.stdout,
        process.stdin,
        abort_transport=abort_transport,
    )
    return process, client


def _stop_helper(process: subprocess.Popen, client: RequestClient) -> None:
    client.close()
    try:
        process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        process.kill()
        process.wait(timeout=5)
    finally:
        if process.stdin is not None:
            process.stdin.close()
        if process.stdout is not None:
            process.stdout.close()


def bench_single_reads(client: RequestClient, path: str) -> None:
    opened = client.call(
        proto.OP_OPEN, {"target": "host", "container": None, "path": path},
    )
    handle_id = opened["handle"]

    started = time.monotonic()
    content, _ = client.call_read(handle_id, 0, 1024 * 1024)
    elapsed_1mib = time.monotonic() - started
    assert len(content) == 1024 * 1024

    started = time.monotonic()
    for offset in range(0, 16 * 1024 * 1024, 4096):
        client.call_read(handle_id, offset, 4096)
    elapsed_4kib_sequence = time.monotonic() - started

    client.call(proto.OP_CLOSE, {"handle": handle_id})

    print(f"1 MiB single read: {elapsed_1mib * 1000:.2f} ms")
    print(
        "4 KiB x4096 sequential reads: "
        f"{elapsed_4kib_sequence:.3f} s total, "
        f"{elapsed_4kib_sequence / 4096 * 1000:.3f} ms/read"
    )


def bench_sequential_throughput(client: RequestClient, path: str) -> float:
    opened = client.call(
        proto.OP_OPEN, {"target": "host", "container": None, "path": path},
    )
    handle_id = opened["handle"]

    started = time.monotonic()
    total_read = 0
    offset = 0
    while offset < THROUGHPUT_FILE_BYTES:
        length = min(THROUGHPUT_CHUNK_BYTES, THROUGHPUT_FILE_BYTES - offset)
        content, _eof = client.call_read(handle_id, offset, length)
        total_read += len(content)
        offset += length
    elapsed = time.monotonic() - started

    client.call(proto.OP_CLOSE, {"handle": handle_id})

    assert total_read == THROUGHPUT_FILE_BYTES
    throughput_mibps = (total_read / (1024 * 1024)) / elapsed
    print(
        f"Sequential throughput: {total_read / (1024 * 1024):.1f} MiB in "
        f"{elapsed * 1000:.2f} ms => {throughput_mibps:.1f} MiB/s"
    )
    return throughput_mibps


def bench_stat_latency(client: RequestClient, path: str) -> float:
    # Warm-up call so the "cold" first stat doesn't skew the loop average;
    # the design's <200ms cold-stat gate is evaluated separately in Phase B
    # against real SSH/container latency, this loopback number is a
    # protocol-overhead floor.
    client.call(proto.OP_STAT, {"path": path})

    started = time.monotonic()
    for _ in range(STAT_CALL_COUNT):
        client.call(proto.OP_STAT, {"path": path})
    elapsed = time.monotonic() - started
    avg_ms = elapsed / STAT_CALL_COUNT * 1000
    print(f"stat: {STAT_CALL_COUNT} calls in {elapsed * 1000:.2f} ms => {avg_ms:.3f} ms/call avg")
    return avg_ms


def bench_list_latency(client: RequestClient, directory: str) -> float:
    started = time.monotonic()
    result = client.call(proto.OP_LIST, {"paths": [directory]})
    elapsed = time.monotonic() - started
    entries = len(result["results"][0]["entries"])
    assert entries == LIST_DIRECTORY_ENTRY_COUNT
    print(
        f"list ({entries} entries): {elapsed * 1000:.3f} ms"
    )
    return elapsed * 1000


def bench_concurrent_reads(client: RequestClient, paths: list[str]) -> float:
    """Prove concurrent READ requests are not serialized behind a single
    global lock: N threads each drive their own handle, all in flight on the
    same RequestClient/connection at once."""
    handles = [
        client.call(
            proto.OP_OPEN, {"target": "host", "container": None, "path": path},
        )["handle"]
        for path in paths
    ]

    errors: list[BaseException] = []

    def read_all(handle_id: int) -> None:
        try:
            offset = 0
            while offset < CONCURRENT_READ_BYTES_PER_READER:
                length = min(
                    CONCURRENT_CHUNK_BYTES, CONCURRENT_READ_BYTES_PER_READER - offset
                )
                content, _eof = client.call_read(handle_id, offset, length)
                assert len(content) == length
                offset += length
        except BaseException as exc:  # noqa: BLE001 - surfaced to the main thread
            errors.append(exc)

    started = time.monotonic()
    threads = [
        threading.Thread(target=read_all, args=(handle,)) for handle in handles
    ]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join(timeout=30)
    elapsed = time.monotonic() - started

    for handle_id in handles:
        client.call(proto.OP_CLOSE, {"handle": handle_id})

    if errors:
        raise errors[0]
    for thread in threads:
        if thread.is_alive():
            raise TimeoutError("a concurrent reader thread did not finish in time")

    total_bytes = CONCURRENT_READERS * CONCURRENT_READ_BYTES_PER_READER
    throughput_mibps = (total_bytes / (1024 * 1024)) / elapsed
    print(
        f"Concurrent reads: {CONCURRENT_READERS} readers x "
        f"{CONCURRENT_READ_BYTES_PER_READER // (1024 * 1024)} MiB in "
        f"{elapsed * 1000:.2f} ms => {throughput_mibps:.1f} MiB/s aggregate"
    )
    return throughput_mibps


def main() -> None:
    if not os.path.exists(HELPER_BINARY):
        raise SystemExit(
            f"dev helper binary not found at {HELPER_BINARY}; build it first with "
            "'go build -o bin/msdev-file-helper-dev .' in tools/msdev/file-helper"
        )

    work_dir = tempfile.mkdtemp(prefix="msdev-bench-loopback-")
    try:
        single_read_path = os.path.join(work_dir, "single_read.bin")
        with open(single_read_path, "wb") as handle:
            handle.write(b"x" * (16 * 1024 * 1024))

        throughput_path = os.path.join(work_dir, "throughput.bin")
        with open(throughput_path, "wb") as handle:
            handle.write(os.urandom(1024 * 1024) * (THROUGHPUT_FILE_BYTES // (1024 * 1024)))

        list_dir = os.path.join(work_dir, "listing")
        os.makedirs(list_dir)
        for index in range(LIST_DIRECTORY_ENTRY_COUNT):
            with open(os.path.join(list_dir, f"entry_{index}.txt"), "w") as handle:
                handle.write(str(index))

        concurrent_paths = []
        for index in range(CONCURRENT_READERS):
            path = os.path.join(work_dir, f"concurrent_{index}.bin")
            with open(path, "wb") as handle:
                handle.write(os.urandom(CONCURRENT_READ_BYTES_PER_READER))
            concurrent_paths.append(path)

        process, client = _start_helper()
        try:
            bench_single_reads(client, single_read_path)
            throughput_mibps = bench_sequential_throughput(client, throughput_path)
            bench_stat_latency(client, single_read_path)
            bench_list_latency(client, list_dir)
            concurrent_mibps = bench_concurrent_reads(client, concurrent_paths)

            print()
            print(f"GATE sequential_throughput_mibps={throughput_mibps:.1f}")
            print(f"GATE concurrent_throughput_mibps={concurrent_mibps:.1f}")
        finally:
            _stop_helper(process, client)
    finally:
        shutil.rmtree(work_dir, ignore_errors=True)


if __name__ == "__main__":
    main()