from __future__ import annotations

import json
import time
from typing import Any

from .benchmark_stats import BenchmarkStats, summarize_runs

_FIELD_ORDER = [
    "lang",
    "scenario",
    "ticks",
    "warmup",
    "iterations",
    "repeat",
    "trace",
    "run",
    "count",
    "runs",
    "min_us",
    "p50_us",
    "p95_us",
    "p99_us",
    "max_us",
    "mean_us",
    "overall_mean_us",
    "total_elapsed_ms",
    "overall_elapsed_ms",
    "throughput_per_sec",
    "best_p50_us",
    "best_p99_us",
    "best_throughput_per_sec",
]


def run(trace_config, benchmark_writer, measured_trace_writer, execute_once) -> None:
    _emit_start(trace_config, benchmark_writer)
    run_stats: list[BenchmarkStats] = []
    for run_ordinal in range(1, trace_config.benchmark_repeat + 1):
        _warmup(trace_config, execute_once)
        stats = _measure(trace_config, measured_trace_writer, execute_once)
        run_stats.append(stats)
        _emit_result(trace_config, benchmark_writer, run_ordinal, stats)
    _emit_summary(trace_config, benchmark_writer, summarize_runs(run_stats))
    benchmark_writer.flush()


def _warmup(trace_config, execute_once) -> None:
    discard_writer = _DiscardingWriter()
    for warmup_ordinal in range(trace_config.benchmark_warmup):
        execute_once(trace_config, discard_writer)


def _measure(trace_config, measured_trace_writer, execute_once) -> BenchmarkStats:
    iteration_durations_ns: list[int] = []
    benchmark_start_ns = time.perf_counter_ns()
    for iteration_ordinal in range(trace_config.benchmark_iterations):
        iteration_start_ns = time.perf_counter_ns()
        execute_once(trace_config, measured_trace_writer)
        iteration_durations_ns.append(time.perf_counter_ns() - iteration_start_ns)
    total_elapsed_ns = time.perf_counter_ns() - benchmark_start_ns
    return BenchmarkStats.from_iteration_nanos(iteration_durations_ns, total_elapsed_ns)


def _emit_start(trace_config, benchmark_writer) -> None:
    event_fields: dict[str, Any] = {
        "lang": trace_config.lang,
        "scenario": trace_config.scenario,
        "ticks": trace_config.ticks,
        "warmup": trace_config.benchmark_warmup,
        "iterations": trace_config.benchmark_iterations,
        "repeat": trace_config.benchmark_repeat,
        "trace": trace_config.trace_enabled,
    }
    benchmark_writer.write(_format_line(trace_config, "benchmark.start", event_fields) + "\n")


def _emit_result(trace_config, benchmark_writer, run_ordinal: int, stats: BenchmarkStats) -> None:
    event_fields: dict[str, Any] = {
        "lang": trace_config.lang,
        "scenario": trace_config.scenario,
        "run": run_ordinal,
        "count": stats.count,
        "min_us": stats.min_us,
        "p50_us": stats.p50_us,
        "p95_us": stats.p95_us,
        "p99_us": stats.p99_us,
        "max_us": stats.max_us,
        "mean_us": stats.mean_us,
        "total_elapsed_ms": stats.total_elapsed_ms,
        "throughput_per_sec": stats.throughput_per_sec,
    }
    benchmark_writer.write(_format_line(trace_config, "benchmark.result", event_fields) + "\n")


def _emit_summary(trace_config, benchmark_writer, summary) -> None:
    event_fields: dict[str, Any] = {
        "lang": trace_config.lang,
        "scenario": trace_config.scenario,
        "runs": summary.runs,
        "overall_mean_us": summary.overall_mean_us,
        "overall_elapsed_ms": summary.overall_elapsed_ms,
        "best_p50_us": summary.best_p50_us,
        "best_p99_us": summary.best_p99_us,
        "best_throughput_per_sec": summary.best_throughput_per_sec,
    }
    benchmark_writer.write(_format_line(trace_config, "benchmark.summary", event_fields) + "\n")


def _format_line(trace_config, event_name: str, event_fields: dict[str, Any]) -> str:
    if trace_config.benchmark_format == "jsonl":
        return _format_json_line(event_name, event_fields)
    return _format_text_line(event_name, event_fields)


def _format_text_line(event_name: str, event_fields: dict[str, Any]) -> str:
    parts = [f"[{event_name}]"]
    for field_name in _ordered_keys(event_fields):
        field_value = event_fields.get(field_name)
        if field_value is None:
            continue
        parts.append(f"{field_name}={_format_scalar(field_value)}")
    return " ".join(parts)


def _format_json_line(event_name: str, event_fields: dict[str, Any]) -> str:
    ordered_payload: dict[str, Any] = {"event": event_name}
    for field_name in _ordered_keys(event_fields):
        field_value = event_fields.get(field_name)
        if field_value is None:
            continue
        ordered_payload[field_name] = _json_value(field_value)
    return json.dumps(ordered_payload, separators=(",", ":"))


def _ordered_keys(event_fields: dict[str, Any]) -> list[str]:
    ordered_keys = [field_name for field_name in _FIELD_ORDER if field_name in event_fields]
    additional_keys = sorted(field_name for field_name in event_fields if field_name not in _FIELD_ORDER)
    return ordered_keys + additional_keys


def _format_scalar(field_value: Any) -> str:
    if isinstance(field_value, float):
        return f"{field_value:.3f}"
    return str(field_value)


def _json_value(field_value: Any) -> Any:
    if isinstance(field_value, float):
        return round(field_value, 3)
    return field_value


class _DiscardingWriter:
    def write(self, payload: str) -> int:
        return len(payload)

    def flush(self) -> None:
        return None
