"""Worker functions and benchmark runners for pywrkr."""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import math
import random
import signal
import ssl
import statistics
import sys
import time
import uuid
from urllib.parse import urlparse

import aiohttp

from pywrkr import reporting as _reporting
from pywrkr.config import (
    ActiveUsers,
    AutofindConfig,
    BenchmarkConfig,
    LatencyBreakdown,
    RequestCounter,
    StepResult,
    WorkerStats,
)
from pywrkr.reporting import (
    _format_latency_short,
    compute_percentiles,
    evaluate_thresholds,
    format_bytes,
    format_duration,
    print_autofind_summary,
    print_results,
    print_threshold_results,
)
from pywrkr.traffic_profiles import RateLimiter

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Module-level constants
# ---------------------------------------------------------------------------

# Verbosity level tags used in log messages
_V3_TAG = "[v3]"
_V4_TAG = "[v4]"

# Progress bar rendering
_PROGRESS_BAR_WIDTH = 20
_PROGRESS_FILLED_CHAR = "\u2588"
_PROGRESS_EMPTY_CHAR = "\u2591"

# Maximum number of unique step names tracked per WorkerStats.
# Prevents unbounded memory growth from dynamic step names.
_MAX_STEP_NAMES = 500


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def _build_progress_bar(current: float, total: float) -> tuple[str, float]:
    """Build a text progress bar and return (bar_string, percentage).

    Args:
        current: Current progress value (e.g. elapsed time or completed requests).
        total: Total target value.

    Returns:
        Tuple of (bar_string, percentage) where bar_string is a rendered
        progress bar of ``_PROGRESS_BAR_WIDTH`` characters and percentage is
        clamped to [0, 100].
    """
    pct = min(current / total * 100, 100.0) if total > 0 else 100.0
    bar_filled = int(pct / 100 * _PROGRESS_BAR_WIDTH)
    bar_empty = _PROGRESS_BAR_WIDTH - bar_filled
    filled = _PROGRESS_FILLED_CHAR * bar_filled
    empty = _PROGRESS_EMPTY_CHAR * bar_empty
    return f"[{filled}{empty}]", pct


def _normalize_config_duration(config: BenchmarkConfig) -> float | None:
    """Return the effective duration from config, normalized early.

    Centralizes the ``config.duration`` check so callers don't need to
    repeat the conditional logic.
    """
    if config.duration is not None and config.duration > 0:
        return config.duration
    return None


def _record_step_latency(stats: WorkerStats, step_name: str, latency: float) -> None:
    """Record a latency sample for a named step, with bounded keys.

    If the number of unique step names exceeds ``_MAX_STEP_NAMES``, new
    step names are folded into a catch-all ``[other steps]`` bucket to
    prevent unbounded memory growth from dynamic step names.
    """
    if step_name in stats.step_latencies:
        stats.step_latencies[step_name].append(latency)
    elif len(stats.step_latencies) < _MAX_STEP_NAMES:
        stats.step_latencies[step_name].append(latency)
    else:
        stats.step_latencies["[other steps]"].append(latency)


# ---------------------------------------------------------------------------
# Live TUI Dashboard
# ---------------------------------------------------------------------------


class LiveDashboard:
    """Real-time terminal dashboard using rich library."""

    def __init__(
        self,
        all_stats: list[WorkerStats],
        config: BenchmarkConfig,
        start_time: float,
        active_users: ActiveUsers | None = None,
    ) -> None:
        """Initialize the dashboard with stats, config, and timing state."""
        self.all_stats = all_stats
        self.config = config
        self.start_time = start_time
        self.active_users = active_users

    def _build_display(self) -> "Panel":  # noqa: F821
        """Build the rich Panel for the current dashboard state."""
        from rich.panel import Panel
        from rich.table import Table

        elapsed = time.monotonic() - self.start_time
        total_req = sum(ws.total_requests for ws in self.all_stats)
        total_err = sum(ws.errors for ws in self.all_stats)
        rps = total_req / elapsed if elapsed > 0 else 0.0
        total_bytes = sum(ws.total_bytes for ws in self.all_stats)
        transfer_rate = total_bytes / elapsed if elapsed > 0 else 0.0
        error_rate = (total_err / total_req * 100) if total_req > 0 else 0.0

        all_latencies = []
        for ws in self.all_stats:
            all_latencies.extend(ws.latencies)

        status_codes: dict[int, int] = {}
        for ws in self.all_stats:
            for code, count in ws.status_codes.items():
                status_codes[code] = status_codes.get(code, 0) + count

        if self.config.users:
            mode_str = f"{self.config.users} users, {self.config.duration}s duration"
        elif self.config.num_requests:
            mode_str = f"{self.config.num_requests} requests"
        else:
            mode_str = f"{self.config.duration}s duration"

        duration = _normalize_config_duration(self.config)
        if duration is not None:
            bar, pct = _build_progress_bar(elapsed, duration)
            progress_str = f"Elapsed: {elapsed:.1f}s / {duration:.1f}s  {bar} {pct:.1f}%"
        elif self.config.num_requests:
            total_n = self.config.num_requests
            bar, pct = _build_progress_bar(total_req, total_n)
            progress_str = f"Progress: {total_req}/{total_n}  {bar} {pct:.1f}%"
        else:
            progress_str = f"Elapsed: {elapsed:.1f}s"

        table = Table(show_header=False, show_edge=False, pad_edge=False, box=None)
        table.add_column("key", style="bold cyan", width=16)
        table.add_column("value")

        table.add_row("Target:", self.config.url)
        table.add_row("Mode:", mode_str)
        table.add_row("", progress_str)
        table.add_row("", "")

        if self.active_users is not None:
            table.add_row("Active Users:", f"{self.active_users.count}")

        table.add_row("Requests:", f"{total_req:,}")
        table.add_row("Errors:", f"{total_err:,} ({error_rate:.1f}%)")
        table.add_row("RPS:", f"{rps:,.1f}")
        table.add_row("Transfer:", f"{format_bytes(transfer_rate)}/s")
        table.add_row("", "")

        if all_latencies:
            pairs = compute_percentiles(all_latencies)
            pair_dict = dict(pairs)
            p50 = format_duration(pair_dict.get(50, 0))
            p95 = format_duration(pair_dict.get(95, 0))
            p99 = format_duration(pair_dict.get(99, 0))
            table.add_row("Latency", f"p50: {p50}  p95: {p95}  p99: {p99}")
        else:
            table.add_row("Latency", "(no data yet)")

        table.add_row("", "")

        if status_codes:
            codes_str = "  ".join(
                f"{code}: {count}" for code, count in sorted(status_codes.items())
            )
            table.add_row("Status Codes:", codes_str)

        max_bar = 24
        if rps > 0:
            bar_len = min(max_bar, max(1, int(rps / max(rps, 1) * max_bar)))
            bar = "\u2588" * bar_len + "\u2591" * (max_bar - bar_len)
            table.add_row("Throughput:", f"{bar} {rps:.0f} req/s")

        return Panel(table, title="pywrkr Live Dashboard", border_style="green")

    async def run(self, stop_event: asyncio.Event) -> None:
        """Update the dashboard every 0.5s until stop_event is set."""
        from rich.live import Live

        with Live(self._build_display(), refresh_per_second=2) as live:
            while not stop_event.is_set():
                try:
                    await asyncio.wait_for(stop_event.wait(), timeout=0.5)
                    break
                except asyncio.TimeoutError:
                    pass
                live.update(self._build_display())


# ---------------------------------------------------------------------------
# URL helpers
# ---------------------------------------------------------------------------


def make_url(url: str, random_param: bool) -> str:
    """Return the URL, optionally appending a unique cache-busting query parameter."""
    if not random_param:
        return url
    separator = "&" if "?" in url else "?"
    return f"{url}{separator}_cb={uuid.uuid4().hex}"


def _build_request_headers(config: BenchmarkConfig) -> dict[str, str]:
    """Build the common request headers from benchmark config.

    Assembles headers from config.headers, adds Basic auth and cookie
    headers if configured. Returns a new dict each call to avoid
    shared mutable state between workers.
    """
    headers = dict(config.headers)
    if config.basic_auth:
        encoded = base64.b64encode(config.basic_auth.encode()).decode()
        headers["Authorization"] = f"Basic {encoded}"
    if config.cookies:
        headers["Cookie"] = "; ".join(config.cookies)
    return headers


def _merge_all_stats(all_stats: list[WorkerStats]) -> WorkerStats:
    """Merge a list of WorkerStats into a single aggregated WorkerStats.

    Combines all counters, latencies, timelines, and breakdowns from
    multiple workers into one unified stats object.

    Step latencies are bounded: unique step names are capped at
    ``_MAX_STEP_NAMES`` to prevent unbounded memory growth.
    """
    merged = WorkerStats()
    for ws in all_stats:
        merged.total_requests += ws.total_requests
        merged.total_bytes += ws.total_bytes
        merged.errors += ws.errors
        merged.content_length_errors += ws.content_length_errors
        merged.latencies.extend(ws.latencies)
        merged.rps_timeline.extend(ws.rps_timeline)
        for k, v in ws.error_types.items():
            merged.error_types[k] += v
        for k, v in ws.status_codes.items():
            merged.status_codes[k] += v
        for k, v in ws.step_latencies.items():
            # Enforce the cap on unique step names during merge
            if k not in merged.step_latencies:
                if len(merged.step_latencies) >= _MAX_STEP_NAMES:
                    k = "[other steps]"
                else:
                    merged.step_latencies[k] = []
            merged.step_latencies[k].extend(v)
        merged.breakdowns.extend(ws.breakdowns)

    logger.debug(
        "Merged stats from %d workers: %d total requests, %d latency samples, %d breakdown samples",
        len(all_stats),
        merged.total_requests,
        len(merged.latencies),
        len(merged.breakdowns),
    )
    return merged


def _create_ssl_context(config: BenchmarkConfig) -> "ssl.SSLContext | None":
    """Create an SSL context based on the benchmark configuration.

    Returns None for plain HTTP. For HTTPS, creates a context with
    certificate verification controlled by config.ssl_config.
    """
    from urllib.parse import urlparse

    parsed = urlparse(config.url)
    if parsed.scheme != "https":
        return None

    ssl_ctx = ssl.create_default_context()
    if not config.ssl_config.verify:
        ssl_ctx.check_hostname = False
        ssl_ctx.verify_mode = ssl.CERT_NONE
    elif config.ssl_config.ca_bundle:
        ssl_ctx.load_verify_locations(config.ssl_config.ca_bundle)
    return ssl_ctx


# ---------------------------------------------------------------------------
# Latency breakdown tracing
# ---------------------------------------------------------------------------


def create_trace_config(stats: WorkerStats) -> aiohttp.TraceConfig:
    """Create an aiohttp TraceConfig that captures per-request latency breakdown.

    The trace context (a dict) stores timing data per request. When the request
    ends, a LatencyBreakdown is computed and appended to stats.breakdowns.
    """
    trace_config = aiohttp.TraceConfig()

    async def on_request_start(session, trace_ctx, params):
        ctx = trace_ctx.trace_request_ctx
        ctx["request_start"] = time.monotonic()
        ctx["dns_start"] = None
        ctx["dns_end"] = None
        ctx["conn_start"] = None
        ctx["conn_end"] = None
        ctx["headers_sent"] = None
        ctx["first_byte"] = None
        ctx["is_reused"] = True  # assume reused; set to False if we see connection creation

    async def on_dns_resolvehost_start(session, trace_ctx, params):
        ctx = trace_ctx.trace_request_ctx
        ctx["dns_start"] = time.monotonic()
        ctx["is_reused"] = False

    async def on_dns_resolvehost_end(session, trace_ctx, params):
        ctx = trace_ctx.trace_request_ctx
        ctx["dns_end"] = time.monotonic()

    async def on_connection_create_start(session, trace_ctx, params):
        ctx = trace_ctx.trace_request_ctx
        ctx["conn_start"] = time.monotonic()
        ctx["is_reused"] = False

    async def on_connection_create_end(session, trace_ctx, params):
        ctx = trace_ctx.trace_request_ctx
        ctx["conn_end"] = time.monotonic()

    async def on_request_headers_sent(session, trace_ctx, params):
        ctx = trace_ctx.trace_request_ctx
        ctx["headers_sent"] = time.monotonic()

    async def on_response_chunk_received(session, trace_ctx, params):
        ctx = trace_ctx.trace_request_ctx
        if ctx.get("first_byte") is None:
            ctx["first_byte"] = time.monotonic()

    async def on_request_end(session, trace_ctx, params):
        ctx = trace_ctx.trace_request_ctx
        end = time.monotonic()

        # Defensive: if on_request_start never fired, skip breakdown
        if ctx.get("request_start") is None:
            logger.debug("Trace context missing 'request_start'; skipping breakdown")
            return

        # DNS time
        dns = 0.0
        if ctx.get("dns_start") is not None and ctx.get("dns_end") is not None:
            dns = ctx["dns_end"] - ctx["dns_start"]

        # TCP connect time (includes TLS if HTTPS)
        connect_total = 0.0
        if ctx.get("conn_start") is not None and ctx.get("conn_end") is not None:
            connect_total = ctx["conn_end"] - ctx["conn_start"]

        # NOTE: TLS time is an *estimate*.  aiohttp's connection_create
        # callback spans TCP+TLS combined and does not provide a separate
        # TLS-only signal.  The value reported here (currently 0.0) is a
        # best-effort approximation; treat the ``tls`` field in
        # LatencyBreakdown as ``tls_estimated`` in any analysis or reports.
        tls = 0.0
        tcp_connect = connect_total  # default: entire connect time is TCP

        # TTFB: from headers_sent to first byte received
        ttfb = 0.0
        headers_sent = ctx.get("headers_sent")
        first_byte = ctx.get("first_byte")
        if headers_sent is not None and first_byte is not None:
            ttfb = first_byte - headers_sent
        elif headers_sent is not None:
            # No chunks received (empty body) -- use end time
            ttfb = end - headers_sent

        # Transfer time: from first byte to end
        transfer = 0.0
        if first_byte is not None:
            transfer = end - first_byte

        bd = LatencyBreakdown(
            dns=max(dns, 0.0),
            connect=max(tcp_connect, 0.0),
            tls=max(tls, 0.0),
            ttfb=max(ttfb, 0.0),
            transfer=max(transfer, 0.0),
            is_reused=ctx.get("is_reused", True),
        )
        stats.breakdowns.append(bd)

    trace_config.on_request_start.append(on_request_start)
    trace_config.on_dns_resolvehost_start.append(on_dns_resolvehost_start)
    trace_config.on_dns_resolvehost_end.append(on_dns_resolvehost_end)
    trace_config.on_connection_create_start.append(on_connection_create_start)
    trace_config.on_connection_create_end.append(on_connection_create_end)
    trace_config.on_request_headers_sent.append(on_request_headers_sent)
    trace_config.on_response_chunk_received.append(on_response_chunk_received)
    trace_config.on_request_end.append(on_request_end)

    return trace_config


def aggregate_breakdowns(breakdowns: list[LatencyBreakdown]) -> dict:
    """Compute aggregate statistics for a list of LatencyBreakdown objects.

    Returns a dict with keys: dns, connect, tls, ttfb, transfer, total.
    Each has: avg, min, max, p50, p95, count.
    Also includes: new_connections, reused_connections.
    """
    if not breakdowns:
        return {}

    new_conns = sum(1 for b in breakdowns if not b.is_reused)
    reused_conns = sum(1 for b in breakdowns if b.is_reused)

    phases = {
        "dns": [b.dns for b in breakdowns],
        "connect": [b.connect for b in breakdowns],
        "tls": [b.tls for b in breakdowns],
        "ttfb": [b.ttfb for b in breakdowns],
        "transfer": [b.transfer for b in breakdowns],
        "total": [b.dns + b.connect + b.tls + b.ttfb + b.transfer for b in breakdowns],
    }

    result = {
        "new_connections": new_conns,
        "reused_connections": reused_conns,
    }

    for name, values in phases.items():
        if not values:
            continue
        sorted_vals = sorted(values)
        n = len(sorted_vals)
        p50_idx = min(int(math.ceil(50 / 100 * n)) - 1, n - 1)
        p95_idx = min(int(math.ceil(95 / 100 * n)) - 1, n - 1)
        result[name] = {
            "avg": statistics.mean(values),
            "min": min(values),
            "max": max(values),
            "p50": sorted_vals[p50_idx],
            "p95": sorted_vals[p95_idx],
            "count": n,
        }

    return result


# ---------------------------------------------------------------------------
# Shared request execution helper
# ---------------------------------------------------------------------------


class _RequestResult:
    """Result from _execute_request; avoids creating dataclass per request."""

    __slots__ = ("latency", "status", "data_len", "error_name", "cancelled")

    def __init__(self) -> None:
        self.latency: float = 0.0
        self.status: int = 0
        self.data_len: int = 0
        self.error_name: str | None = None
        self.cancelled: bool = False


async def _execute_request(
    session: aiohttp.ClientSession,
    method: str,
    url: str,
    headers: dict[str, str],
    body: bytes | None,
    ssl_verify: bool,
    timeout: aiohttp.ClientTimeout,
    stats: WorkerStats,
    config: BenchmarkConfig,
    trace_ctx: dict[str, object] | None,
    expected_length_ref: list[int | None],
    step_name: str | None = None,
    assert_status: int | None = None,
    assert_body_contains: str | None = None,
    log_prefix: str = "",
) -> _RequestResult:
    """Execute a single HTTP request and record stats.

    This is the shared core extracted from worker(), user_worker(), and
    scenario_worker(). Handles: request execution, latency recording, status
    code counting, content-length verification, error handling, and step
    latency tracking.

    Returns a _RequestResult with outcome details.
    """
    result = _RequestResult()
    req_start = time.monotonic()
    try:
        async with session.request(
            method,
            url,
            headers=headers,
            data=body,
            ssl=ssl_verify,
            timeout=timeout,
            trace_request_ctx=trace_ctx,
        ) as resp:
            data = await resp.read()
            latency = time.monotonic() - req_start
            result.latency = latency
            result.status = resp.status
            result.data_len = len(data)

            stats.total_requests += 1
            stats.total_bytes += len(data)
            stats.latencies.append(latency)
            stats.status_codes[resp.status] += 1
            if step_name:
                _record_step_latency(stats, step_name, latency)

            # Content-length verification (ab -l style)
            if config.verify_content_length:
                cl = resp.headers.get("Content-Length")
                if cl is not None:
                    declared = int(cl)
                    if expected_length_ref[0] is None:
                        expected_length_ref[0] = declared
                    if declared != expected_length_ref[0] or len(data) != declared:
                        stats.content_length_errors += 1

            # Assertion checks (scenario mode)
            assertion_failed = False
            if assert_status is not None and resp.status != assert_status:
                stats.errors += 1
                err_msg = f"AssertStatus: expected {assert_status}, got {resp.status}"
                stats.error_types[err_msg] += 1
                assertion_failed = True

            if assert_body_contains is not None:
                body_text = data.decode("utf-8", errors="replace")
                if assert_body_contains not in body_text:
                    stats.errors += 1
                    err_msg = f"AssertBody: '{assert_body_contains}' not found"
                    stats.error_types[err_msg] += 1
                    assertion_failed = True

            if not assertion_failed and resp.status >= 400:
                stats.errors += 1
                stats.error_types[f"HTTP {resp.status}"] += 1

            if config.verbosity >= 4:
                logger.debug(
                    "%s %s %s -> %s (%dB, %s)",
                    _V4_TAG,
                    method,
                    url,
                    resp.status,
                    len(data),
                    format_duration(latency),
                )
            elif config.verbosity >= 3:
                logger.debug("%s %s", _V3_TAG, resp.status)

    except asyncio.CancelledError:
        result.cancelled = True
    except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as e:
        latency = time.monotonic() - req_start
        result.latency = latency
        error_name = type(e).__name__
        result.error_name = error_name
        stats.total_requests += 1
        stats.errors += 1
        stats.error_types[error_name] += 1
        stats.latencies.append(latency)
        if step_name:
            stats.step_latencies[step_name].append(latency)
        logger.warning("%sRequest error: %s: %s", log_prefix, error_name, e)

    return result


def _build_session_kwargs(
    connector: aiohttp.TCPConnector,
    config: BenchmarkConfig,
    stats: WorkerStats,
) -> dict:
    """Build kwargs for aiohttp.ClientSession including optional trace config."""
    kwargs: dict = {"connector": connector, "connector_owner": False}
    if config.latency_breakdown:
        kwargs["trace_configs"] = [create_trace_config(stats)]
    return kwargs


async def _think_time_wait(
    think: float,
    jitter: float,
    stop_event: asyncio.Event,
) -> bool:
    """Sleep for think time with jitter. Returns True if stop_event was set."""
    if think <= 0 or stop_event.is_set():
        return stop_event.is_set()
    lo = think * (1 - jitter)
    hi = think * (1 + jitter)
    delay = random.uniform(lo, hi)
    try:
        await asyncio.wait_for(stop_event.wait(), timeout=delay)
        return True  # stop_event was set
    except asyncio.TimeoutError:
        return False  # think time elapsed normally


def _calc_effective_timeout(
    config: BenchmarkConfig,
    start_time: float,
) -> float:
    """Calculate effective request timeout considering remaining duration.

    Returns at least 0.1s to avoid zero/negative timeouts when the
    benchmark has overrun its duration window.
    """
    if config.duration is not None:
        remaining = config.duration - (time.monotonic() - start_time)
        return max(0.1, min(config.timeout_sec, remaining + 1))
    return config.timeout_sec


# ---------------------------------------------------------------------------
# Workers
# ---------------------------------------------------------------------------


async def worker(
    config: BenchmarkConfig,
    stats: WorkerStats,
    connector: aiohttp.TCPConnector,
    stop_event: asyncio.Event,
    request_counter: RequestCounter | None = None,
    rate_limiter: RateLimiter | None = None,
) -> None:
    """Async worker coroutine that sends HTTP requests in a loop.

    Executes HTTP requests against the configured URL until the stop condition
    is met (duration elapsed, request count reached, or stop_event set).
    """
    logger.debug("Worker starting (target=%s)", config.url)
    start_time = time.monotonic()
    interval_start = start_time
    interval_count = 0

    req_headers = _build_request_headers(config)
    expected_length_ref: list[int | None] = [None]
    session_kwargs = _build_session_kwargs(connector, config, stats)

    async with aiohttp.ClientSession(**session_kwargs) as session:
        while not stop_event.is_set():
            if config.duration is not None:
                elapsed = time.monotonic() - start_time
                if elapsed >= config.duration:
                    break

            if request_counter is not None:
                if request_counter.remaining <= 0:
                    break
                request_counter.remaining -= 1

            if rate_limiter is not None:
                await rate_limiter.acquire()
                if stop_event.is_set():
                    break

            effective_timeout = _calc_effective_timeout(config, start_time)
            client_timeout = aiohttp.ClientTimeout(total=effective_timeout)
            request_url = make_url(config.url, config.random_param)
            trace_ctx = {} if config.latency_breakdown else None

            result = await _execute_request(
                session,
                config.method,
                request_url,
                req_headers,
                config.body,
                config.ssl_config.verify,
                client_timeout,
                stats,
                config,
                trace_ctx,
                expected_length_ref,
            )
            if result.cancelled:
                break

            interval_count += 1
            now = time.monotonic()
            if now - interval_start >= 1.0:
                stats.rps_timeline.append((interval_start, interval_count))
                interval_start = now
                interval_count = 0

        if interval_count > 0:
            stats.rps_timeline.append((interval_start, interval_count))

    logger.debug(
        "Worker finished: %d requests, %d errors",
        stats.total_requests,
        stats.errors,
    )


async def user_worker(
    user_id: int,
    config: BenchmarkConfig,
    stats: WorkerStats,
    connector: aiohttp.TCPConnector,
    stop_event: asyncio.Event,
    start_time: float,
    active_users: ActiveUsers,
    rate_limiter: RateLimiter | None = None,
) -> None:
    """Simulate a single virtual user with configurable think time."""
    logger.debug("User %d starting", user_id)
    req_headers = _build_request_headers(config)
    expected_length_ref: list[int | None] = [None]
    active_users.count += 1
    session_kwargs = _build_session_kwargs(connector, config, stats)

    try:
        async with aiohttp.ClientSession(**session_kwargs) as session:
            while not stop_event.is_set():
                elapsed = time.monotonic() - start_time
                if config.duration is not None and elapsed >= config.duration:
                    break

                if rate_limiter is not None and config.think_time == 0:
                    await rate_limiter.acquire()
                    if stop_event.is_set():
                        break

                effective_timeout = _calc_effective_timeout(config, start_time)
                client_timeout = aiohttp.ClientTimeout(total=effective_timeout)
                request_url = make_url(config.url, config.random_param)
                trace_ctx = {} if config.latency_breakdown else None

                result = await _execute_request(
                    session,
                    config.method,
                    request_url,
                    req_headers,
                    config.body,
                    config.ssl_config.verify,
                    client_timeout,
                    stats,
                    config,
                    trace_ctx,
                    expected_length_ref,
                    log_prefix=f"User {user_id} ",
                )
                if result.cancelled:
                    break

                now = time.monotonic()
                stats.rps_timeline.append((now, 1))

                if await _think_time_wait(config.think_time, config.think_time_jitter, stop_event):
                    break
    finally:
        active_users.count -= 1
        logger.debug(
            "User %d finished: %d requests, %d errors",
            user_id,
            stats.total_requests,
            stats.errors,
        )


def _prepare_step_body(step_body, headers: dict) -> bytes | None:
    """Serialize a scenario step body and set Content-Type if needed."""
    if step_body is None:
        return None
    if isinstance(step_body, dict):
        if "Content-Type" not in headers:
            headers["Content-Type"] = "application/json"
        return json.dumps(step_body).encode()
    if isinstance(step_body, str):
        return step_body.encode()
    return step_body


async def scenario_worker(
    user_id: int,
    config: BenchmarkConfig,
    stats: WorkerStats,
    connector: aiohttp.TCPConnector,
    stop_event: asyncio.Event,
    start_time: float,
    active_users: ActiveUsers,
    request_counter: RequestCounter | None = None,
) -> None:
    """Execute a scripted multi-step scenario in a loop."""
    logger.debug("Scenario user %d starting", user_id)
    scenario = config.scenario
    if not scenario:
        return

    base_headers = _build_request_headers(config)
    parsed = urlparse(config.url)
    base_url = f"{parsed.scheme}://{parsed.netloc}"
    expected_length_ref: list[int | None] = [None]

    active_users.count += 1
    try:
        session_kwargs = _build_session_kwargs(connector, config, stats)
        async with aiohttp.ClientSession(**session_kwargs) as session:
            while not stop_event.is_set():
                for step in scenario.steps:
                    if stop_event.is_set():
                        break

                    if config.duration is not None:
                        elapsed = time.monotonic() - start_time
                        if elapsed >= config.duration:
                            return

                    if request_counter is not None:
                        if request_counter.remaining <= 0:
                            return
                        request_counter.remaining -= 1

                    effective_timeout = _calc_effective_timeout(config, start_time)
                    client_timeout = aiohttp.ClientTimeout(total=effective_timeout)
                    request_url = make_url(f"{base_url}{step.path}", config.random_param)

                    req_headers = dict(base_headers)
                    req_headers.update(step.headers)
                    body = _prepare_step_body(step.body, req_headers)

                    step_name = step.name or f"{step.method} {step.path}"

                    result = await _execute_request(
                        session,
                        step.method,
                        request_url,
                        req_headers,
                        body,
                        config.ssl_config.verify,
                        client_timeout,
                        stats,
                        config,
                        None,
                        expected_length_ref,
                        step_name=step_name,
                        assert_status=step.assert_status,
                        assert_body_contains=step.assert_body_contains,
                        log_prefix=f"Scenario user {user_id} step '{step_name}' ",
                    )
                    if result.cancelled:
                        return

                    now = time.monotonic()
                    stats.rps_timeline.append((now, 1))

                    think = step.think_time if step.think_time is not None else scenario.think_time
                    if think <= 0 and config.think_time > 0:
                        think = config.think_time
                    if await _think_time_wait(think, config.think_time_jitter, stop_event):
                        return
    finally:
        active_users.count -= 1
        logger.debug(
            "Scenario user %d finished: %d requests, %d errors",
            user_id,
            stats.total_requests,
            stats.errors,
        )


async def show_progress(
    start: float,
    duration: float | None,
    total_requests: int | None,
    all_stats: list[WorkerStats],
    stop: asyncio.Event,
    active_users: ActiveUsers | None = None,
) -> None:
    """Display a text-based progress line during benchmark execution."""
    while not stop.is_set():
        await asyncio.sleep(1)
        elapsed = time.monotonic() - start
        total_req = sum(ws.total_requests for ws in all_stats)
        total_err = sum(ws.errors for ws in all_stats)
        rps = total_req / elapsed if elapsed > 0 else 0

        users_str = ""
        if active_users is not None:
            users_str = f" | {active_users.count:>5} users"

        if duration is not None:
            _, pct = _build_progress_bar(elapsed, duration)
            sys.stdout.write(
                f"\r  [{pct:5.1f}%] {total_req:>8} requests "
                f"| {rps:>8.1f} req/s | {total_err} errors{users_str} "
            )
        elif total_requests is not None:
            _, pct = _build_progress_bar(total_req, total_requests)
            sys.stdout.write(
                f"\r  [{pct:5.1f}%] {total_req:>8}/{total_requests} requests "
                f"| {rps:>8.1f} req/s | {total_err} errors{users_str} "
            )
        sys.stdout.flush()
    sys.stdout.write("\r" + " " * 80 + "\r")
    sys.stdout.flush()


# ---------------------------------------------------------------------------
# Shared runner helpers
# ---------------------------------------------------------------------------


def _setup_signal_handlers(stop_event: asyncio.Event) -> None:
    """Register SIGINT/SIGTERM handlers that set the stop event."""
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, stop_event.set)


def _create_rate_limiter(config: BenchmarkConfig, duration: float | None) -> RateLimiter | None:
    """Create a rate limiter from config, or None if rate limiting is disabled."""
    if config.rate is None:
        return None
    ramp_duration = duration if config.rate_ramp is not None else None
    return RateLimiter(
        rate=config.rate,
        end_rate=config.rate_ramp,
        ramp_duration=ramp_duration,
        traffic_profile=config.traffic_profile,
        duration=duration,
    )


def _create_progress_task(
    config: BenchmarkConfig,
    all_stats: list[WorkerStats],
    start_time: float,
    stop_event: asyncio.Event,
    *,
    duration: float | None = None,
    num_requests: int | None = None,
    active_users: ActiveUsers | None = None,
    quiet: bool = False,
) -> asyncio.Task:
    """Create the progress display or live dashboard task."""
    if quiet:

        async def _wait_stop(stop):
            await stop.wait()

        return asyncio.create_task(_wait_stop(stop_event))

    if config.live_dashboard and _reporting.RICH_AVAILABLE:
        dashboard = LiveDashboard(all_stats, config, start_time, active_users)
        return asyncio.create_task(dashboard.run(stop_event))

    if config.live_dashboard and not _reporting.RICH_AVAILABLE:
        logger.warning("--live requires 'rich' package. Install with: pip install pywrkr[tui]")
        logger.warning("Falling back to standard progress display.")

    return asyncio.create_task(
        show_progress(start_time, duration, num_requests, all_stats, stop_event, active_users)
    )


async def _finalize_run(
    tasks: list[asyncio.Task],
    stop_event: asyncio.Event,
    progress_task: asyncio.Task,
    connector: aiohttp.TCPConnector,
    all_stats: list[WorkerStats],
    start_time: float,
    config: BenchmarkConfig,
    rate_limiter: RateLimiter | None,
    concurrency: int,
    *,
    quiet: bool = False,
) -> tuple[WorkerStats, int]:
    """Await workers, merge stats, print results, and evaluate thresholds."""
    try:
        await asyncio.gather(*tasks, return_exceptions=True)
        # Signal all workers to stop before accessing shared stats to avoid
        # race conditions where a still-running worker mutates stats during merge.
        stop_event.set()
        await progress_task
    finally:
        await connector.close()

    end_time = time.monotonic()
    actual_duration = end_time - start_time
    merged = _merge_all_stats(all_stats)

    if not quiet:
        print_results(merged, actual_duration, concurrency, start_time, config, rate_limiter)

    exit_code = 0
    if config.thresholds:
        th_results = evaluate_thresholds(config.thresholds, merged, actual_duration)
        if not quiet:
            print_threshold_results(th_results, file=sys.stdout)
        if any(not passed for _, _, passed in th_results):
            exit_code = 2

    return merged, exit_code


# ---------------------------------------------------------------------------
# Main benchmark runner
# ---------------------------------------------------------------------------


async def run_benchmark(config: BenchmarkConfig) -> tuple[WorkerStats, int]:
    """Run a fixed-concurrency benchmark and return merged stats with exit code.

    Creates N worker tasks distributed across thread groups, each sharing a
    connection pool via aiohttp.TCPConnector. Supports duration-based and
    request-count modes, with optional rate limiting and live dashboard.

    Args:
        config: Full benchmark configuration including URL, concurrency,
            duration/request count, and output options.

    Returns:
        Tuple of (merged_stats, exit_code) where exit_code is 0 for success
        or 2 if any SLO threshold was breached.

    Raises:
        No exceptions are raised to the caller. Network errors, timeouts,
        and HTTP errors are captured in WorkerStats.

    Concurrency notes:
        Workers are distributed across thread groups. Each group shares a
        TCPConnector with a connection limit equal to the number of workers
        in that group. Signal handlers (SIGINT/SIGTERM) set a stop_event
        that all workers check between requests.
    """
    mode_str = (
        f"{config.num_requests} requests" if config.num_requests else f"{config.duration}s duration"
    )
    logger.info("Running benchmark: %s", config.url)
    logger.info(
        "  %d worker groups, %d connections, %s", config.threads, config.connections, mode_str
    )
    logger.info(
        "  Method: %s, Timeout: %ss, Keep-Alive: %s",
        config.method,
        config.timeout_sec,
        "yes" if config.keepalive else "no",
    )
    if config.rate is not None:
        rate_str = f"{config.rate:,.0f} req/s"
        if config.rate_ramp is not None:
            rate_str += f" -> {config.rate_ramp:,.0f} req/s (ramp)"
        logger.info("  Rate Limit: %s", rate_str)
    if config.random_param:
        logger.info("  Cache-Buster: random _cb= parameter per request")
    if config.basic_auth:
        logger.info("  Auth: Basic (user=%s)", config.basic_auth.split(":")[0])
    if config.cookies:
        logger.info("  Cookies: %d", len(config.cookies))
    logger.info("")

    stop_event = asyncio.Event()
    _setup_signal_handlers(stop_event)

    rate_limiter = _create_rate_limiter(config, config.duration)

    # Distribute connections across worker groups
    conns_per_group = max(1, config.connections // config.threads)
    remainder = config.connections % config.threads

    # Shared counter for request-count mode
    request_counter: RequestCounter | None = None
    if config.num_requests is not None:
        request_counter = RequestCounter(config.num_requests)

    all_stats: list[WorkerStats] = []
    tasks = []
    start_time = time.monotonic()
    ssl_ctx = _create_ssl_context(config)

    connector = aiohttp.TCPConnector(
        limit=config.connections,
        ssl=ssl_ctx,
        force_close=not config.keepalive,
        enable_cleanup_closed=True,
    )

    # Log actual worker distribution across groups
    group_sizes = []
    for i in range(config.threads):
        n_conns = conns_per_group + (1 if i < remainder else 0)
        if n_conns == 0:
            continue
        group_sizes.append(n_conns)

        for j in range(n_conns):
            ws = WorkerStats()
            all_stats.append(ws)
            if config.scenario:
                _active = ActiveUsers()
                tasks.append(
                    asyncio.create_task(
                        scenario_worker(
                            j,
                            config,
                            ws,
                            connector,
                            stop_event,
                            start_time,
                            _active,
                            request_counter,
                        )
                    )
                )
            else:
                tasks.append(
                    asyncio.create_task(
                        worker(config, ws, connector, stop_event, request_counter, rate_limiter)
                    )
                )

    logger.debug(
        "Worker distribution: %d groups, sizes=%s, total=%d workers",
        len(group_sizes),
        group_sizes,
        sum(group_sizes),
    )

    progress_task = _create_progress_task(
        config,
        all_stats,
        start_time,
        stop_event,
        duration=config.duration,
        num_requests=config.num_requests,
    )

    return await _finalize_run(
        tasks,
        stop_event,
        progress_task,
        connector,
        all_stats,
        start_time,
        config,
        rate_limiter,
        config.connections,
    )


# ---------------------------------------------------------------------------
# User simulation runner
# ---------------------------------------------------------------------------


async def run_user_simulation(config: BenchmarkConfig) -> tuple[WorkerStats, int]:
    """Run a virtual-user load test with ramp-up and think time.

    Creates one task per virtual user, optionally staggering their start
    times over a ramp-up period. Each user sends requests with configurable
    think time between them.

    Args:
        config: Benchmark configuration. Must have config.users set.
            config.duration defaults to 60s if not specified.

    Returns:
        Tuple of (merged_stats, exit_code) where exit_code is 0 for success
        or 2 if any SLO threshold was breached.

    Concurrency notes:
        All users share a single TCPConnector with
        limit=min(num_users, config.connections). Ramp-up is implemented by
        sleeping between task creation calls, so early users begin sending
        requests while later users are still being launched.
    """
    num_users = config.users
    duration = config.duration or 60.0
    quiet = getattr(config, "_quiet", False)

    if not quiet:
        logger.info("Running user simulation: %s", config.url)
        logger.info("  %d virtual users, %s duration", num_users, format_duration(duration))
        logger.info(
            "  Ramp-up: %s, Think time: %s (jitter: %s)",
            format_duration(config.ramp_up),
            format_duration(config.think_time),
            f"{config.think_time_jitter:.0%}",
        )
        logger.info(
            "  Method: %s, Timeout: %ss, Keep-Alive: %s",
            config.method,
            config.timeout_sec,
            "yes" if config.keepalive else "no",
        )
        if config.rate is not None:
            rate_str = f"{config.rate:,.0f} req/s"
            if config.rate_ramp is not None:
                rate_str += f" -> {config.rate_ramp:,.0f} req/s (ramp)"
            logger.info("  Rate Limit: %s", rate_str)
        if config.random_param:
            logger.info("  Cache-Buster: random _cb= parameter per request")
        logger.info("")

    stop_event = asyncio.Event()
    _setup_signal_handlers(stop_event)

    rate_limiter = _create_rate_limiter(config, duration)

    ssl_ctx = _create_ssl_context(config)

    # Users make sequential requests with think time, so they don't all need
    # a connection simultaneously.  Cap the pool at the configured connections
    # value (defaults to 10) or num_users, whichever is smaller.
    pool_limit = min(num_users, config.connections)
    connector = aiohttp.TCPConnector(
        limit=pool_limit,
        ssl=ssl_ctx,
        force_close=not config.keepalive,
        enable_cleanup_closed=True,
    )
    logger.debug(
        "User simulation: %d users, pool_limit=%d",
        num_users,
        pool_limit,
    )

    all_stats: list[WorkerStats] = []
    tasks = []
    active_users = ActiveUsers()
    start_time = time.monotonic()

    # Ramp-up: stagger user launches
    ramp_delay = config.ramp_up / num_users if config.ramp_up > 0 and num_users > 1 else 0

    progress_task = _create_progress_task(
        config,
        all_stats,
        start_time,
        stop_event,
        duration=duration,
        active_users=active_users,
        quiet=quiet,
    )

    for i in range(num_users):
        if stop_event.is_set():
            break
        ws = WorkerStats()
        all_stats.append(ws)
        if config.scenario:
            tasks.append(
                asyncio.create_task(
                    scenario_worker(i, config, ws, connector, stop_event, start_time, active_users)
                )
            )
        else:
            tasks.append(
                asyncio.create_task(
                    user_worker(
                        i, config, ws, connector, stop_event, start_time, active_users, rate_limiter
                    )
                )
            )
        if ramp_delay > 0 and i < num_users - 1:
            await asyncio.sleep(ramp_delay)

    return await _finalize_run(
        tasks,
        stop_event,
        progress_task,
        connector,
        all_stats,
        start_time,
        config,
        rate_limiter,
        num_users,
        quiet=quiet,
    )


# ---------------------------------------------------------------------------
# Autofind (auto-ramping / step load)
# ---------------------------------------------------------------------------


def _step_passed(step: StepResult, config: AutofindConfig) -> bool:
    """Check whether a step result meets the autofind thresholds."""
    if step.error_rate > config.max_error_rate:
        return False
    if step.p95 > config.max_p95:
        return False
    return True


def _extract_step_result(
    stats: WorkerStats, duration: float, num_users: int, config: AutofindConfig
) -> StepResult:
    """Extract a StepResult from merged WorkerStats."""
    rps = stats.total_requests / duration if duration > 0 else 0.0
    error_rate = (stats.errors / stats.total_requests * 100) if stats.total_requests > 0 else 0.0

    if stats.latencies:
        sorted_lat = sorted(stats.latencies)
        n = len(sorted_lat)
        p50 = sorted_lat[min(int(math.ceil(50 / 100 * n)) - 1, n - 1)]
        p95 = sorted_lat[min(int(math.ceil(95 / 100 * n)) - 1, n - 1)]
        p99 = sorted_lat[min(int(math.ceil(99 / 100 * n)) - 1, n - 1)]
    else:
        p50 = p95 = p99 = 0.0

    result = StepResult(
        users=num_users,
        rps=rps,
        p50=p50,
        p95=p95,
        p99=p99,
        error_rate=error_rate,
        total_requests=stats.total_requests,
        total_errors=stats.errors,
        passed=True,  # will be set below
    )
    result.passed = _step_passed(result, config)
    return result


async def run_autofind(config: AutofindConfig) -> list[StepResult]:
    """Auto-ramp load to find maximum sustainable capacity.

    Starts with start_users, doubles (or multiplies by step_multiplier) each
    step. When a step fails thresholds, binary-searches between the last good
    and first bad user count to refine the answer.
    """
    logger.info("Autofind: ramping load on %s", config.url)
    logger.info(
        "  Thresholds: max error rate=%s%%, max p95=%ss", config.max_error_rate, config.max_p95
    )
    logger.info(
        "  Step duration: %ss, start users: %s, max users: %s",
        config.step_duration,
        config.start_users,
        config.max_users,
    )
    logger.info("  Step multiplier: %sx", config.step_multiplier)
    logger.info("")

    steps: list[StepResult] = []
    last_good: int | None = None
    first_bad: int | None = None
    current_users = config.start_users

    async def _run_step(num_users: int) -> StepResult:
        bench_config = BenchmarkConfig(
            url=config.url,
            users=num_users,
            duration=config.step_duration,
            think_time=config.think_time,
            think_time_jitter=config.think_time_jitter,
            timeout_sec=config.timeout_sec,
            keepalive=config.keepalive,
            random_param=config.random_param,
            ramp_up=0.0,
            _quiet=True,
        )
        stats, _ = await run_user_simulation(bench_config)
        return _extract_step_result(stats, config.step_duration, num_users, config)

    # Phase 1: Exponential ramp-up
    while current_users <= config.max_users:
        logger.info("  Step: testing %d users ...", current_users)
        result = await _run_step(current_users)
        steps.append(result)
        status = "OK" if result.passed else "FAIL"
        logger.info(
            "  %s rps, p95=%s, err=%s%% -> %s",
            f"{result.rps:.1f}",
            _format_latency_short(result.p95),
            f"{result.error_rate:.1f}",
            status,
        )

        if result.passed:
            last_good = current_users
            next_users = int(current_users * config.step_multiplier)
            if next_users == current_users:
                next_users = current_users + 1
            current_users = next_users
        else:
            first_bad = current_users
            break
    else:
        # Reached max_users without failure
        print_autofind_summary(steps, last_good)
        if config.json_output:
            _write_autofind_json(config, steps, last_good)
        return steps

    # Phase 2: Binary search refinement between last_good and first_bad
    if last_good is not None and first_bad is not None and first_bad - last_good > 1:
        lo, hi = last_good, first_bad
        while hi - lo > max(1, lo // 10):  # refine until gap is <10% of lo
            mid = (lo + hi) // 2
            if mid == lo or mid == hi:
                break
            logger.info("  Refine: testing %d users ...", mid)
            result = await _run_step(mid)
            steps.append(result)
            status = "OK" if result.passed else "FAIL"
            logger.info(
                "  %s rps, p95=%s, err=%s%% -> %s",
                f"{result.rps:.1f}",
                _format_latency_short(result.p95),
                f"{result.error_rate:.1f}",
                status,
            )

            if result.passed:
                lo = mid
                last_good = mid
            else:
                hi = mid

    print_autofind_summary(steps, last_good)
    if config.json_output:
        _write_autofind_json(config, steps, last_good)
    return steps


def _write_autofind_json(
    config: AutofindConfig, steps: list[StepResult], max_users: int | None
) -> None:
    """Write autofind results to a JSON file."""
    data = {
        "url": config.url,
        "max_error_rate": config.max_error_rate,
        "max_p95": config.max_p95,
        "step_duration": config.step_duration,
        "max_sustainable_users": max_users,
        "steps": [
            {
                "users": s.users,
                "rps": round(s.rps, 2),
                "p50": round(s.p50, 4),
                "p95": round(s.p95, 4),
                "p99": round(s.p99, 4),
                "error_rate": round(s.error_rate, 2),
                "total_requests": s.total_requests,
                "total_errors": s.total_errors,
                "passed": s.passed,
            }
            for s in steps
        ],
    }
    if config.json_output is None:
        return
    with open(config.json_output, "w") as f:
        json.dump(data, f, indent=2)
    logger.info("  JSON results written to %s", config.json_output)
