import argparse
import asyncio
import json
import math
import os
import socket
import struct
import time
from pathlib import Path

PAYLOAD_MARKER = b"PMSTRESS"


def build_payload(run_id: int, stream_number: int, shard: int, sequence: int, payload_bytes: int) -> bytes:
    prefix = struct.pack("!QQII", run_id, stream_number, shard, sequence)
    padding = bytearray(
        ((stream_number + sequence + i) % 251 for i in range(payload_bytes - len(prefix)))
    )
    if sequence == 0:
        padding[:len(PAYLOAD_MARKER)] = PAYLOAD_MARKER
    return prefix + padding


def percentile(values: list[float], fraction: float) -> float:
    if not values:
        return math.inf
    values.sort()
    return values[max(0, math.ceil(len(values) * fraction) - 1)]


async def send_stream(
    stream_number: int,
    host: str,
    port: int,
    frames: int,
    payload_bytes: int,
    run_id: int,
    shard: int,
) -> float:
    started = time.perf_counter()
    reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=10)
    sock = writer.get_extra_info("socket")
    if sock is not None:
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

    try:
        for sequence in range(frames):
            payload = build_payload(run_id, stream_number, shard, sequence, payload_bytes)
            writer.write(struct.pack("!I", len(payload)) + payload)
            await writer.drain()
            response_size = struct.unpack("!I", await asyncio.wait_for(reader.readexactly(4), 10))[0]
            response = await asyncio.wait_for(reader.readexactly(response_size), 10)
            if response_size != len(payload) or response != payload:
                raise ValueError(
                    f"invalid echo for sequence={sequence}: expected={len(payload)} actual={response_size}"
                )
    finally:
        writer.close()
        await writer.wait_closed()
    return (time.perf_counter() - started) * 1000


async def run(args: argparse.Namespace) -> dict:
    next_stream = 0
    counter_lock = asyncio.Lock()
    latencies: list[float] = []
    errors: list[str] = []
    abort_file = Path(args.abort_file) if args.abort_file else None

    async def worker() -> None:
        nonlocal next_stream
        while True:
            async with counter_lock:
                if abort_file and abort_file.exists():
                    return
                if next_stream >= args.streams:
                    return
                local_number = next_stream
                next_stream += 1
            if args.minimum_duration_seconds > 0:
                scheduled_at = started + (local_number * args.minimum_duration_seconds / args.streams)
                delay = scheduled_at - time.perf_counter()
                if delay > 0:
                    await asyncio.sleep(delay)
            stream_number = args.stream_offset + local_number
            try:
                latency = await send_stream(
                    stream_number,
                    args.host,
                    args.port,
                    args.frames,
                    args.payload_bytes,
                    args.run_id,
                    args.shard,
                )
                latencies.append(latency)
            except Exception as exc:  # individual failures are reported and bounded by the success threshold
                if len(errors) < 100:
                    errors.append(f"stream={stream_number}: {type(exc).__name__}: {exc}")

    started_at_unix_seconds = time.time()
    started = time.perf_counter()
    await asyncio.gather(*(worker() for _ in range(args.concurrency)))
    elapsed = time.perf_counter() - started
    finished_at_unix_seconds = time.time()
    successes = len(latencies)
    result = {
        "runId": args.run_id,
        "shard": args.shard,
        "requestedStreams": args.streams,
        "successfulStreams": successes,
        "failedStreams": args.streams - successes,
        "framesPerStream": args.frames,
        "payloadBytes": args.payload_bytes,
        "minimumDurationSeconds": args.minimum_duration_seconds,
        "startedAtUnixSeconds": round(started_at_unix_seconds, 6),
        "finishedAtUnixSeconds": round(finished_at_unix_seconds, 6),
        "elapsedSeconds": round(elapsed, 3),
        "streamsPerSecond": round(successes / elapsed, 3) if elapsed else 0,
        "latencyP50Millis": round(percentile(latencies, 0.50), 3) if latencies else None,
        "latencyP95Millis": round(percentile(latencies, 0.95), 3) if latencies else None,
        "latencyP99Millis": round(percentile(latencies, 0.99), 3) if latencies else None,
        "latencyMaxMillis": round(max(latencies, default=0), 3),
        "abortedByCaptureMonitor": bool(abort_file and abort_file.exists()),
        "sampleErrors": errors,
    }
    Path(args.report).write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result), flush=True)
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate real framed TCP traffic for Packmate")
    parser.add_argument("--host", default=os.getenv("STRESS_TARGET_HOST", "172.31.250.10"))
    parser.add_argument("--port", type=int, default=int(os.getenv("STRESS_TARGET_PORT", "18080")))
    parser.add_argument("--streams", type=int, required=True)
    parser.add_argument("--frames", type=int, default=8)
    parser.add_argument("--payload-bytes", type=int, default=1024)
    parser.add_argument("--concurrency", type=int, default=256)
    parser.add_argument("--run-id", type=int, required=True)
    parser.add_argument("--shard", type=int, required=True)
    parser.add_argument("--stream-offset", type=int, required=True)
    parser.add_argument("--report", required=True)
    parser.add_argument("--abort-file")
    parser.add_argument("--minimum-duration-seconds", type=float, default=0)
    args = parser.parse_args()
    if (
        args.streams <= 0
        or args.frames <= 0
        or args.payload_bytes < 32
        or args.concurrency <= 0
        or args.minimum_duration_seconds < 0
    ):
        parser.error("streams, frames and concurrency must be positive; payload must be at least 32 bytes")
    return args


if __name__ == "__main__":
    arguments = parse_args()
    outcome = asyncio.run(run(arguments))
    raise SystemExit(0 if outcome["successfulStreams"] == arguments.streams else 1)