from __future__ import annotations

import argparse
import json
import platform
import time
from typing import Any

from kernite.evaluator import evaluate_execute


def _percentile(sorted_values: list[float], percentile: float) -> float:
    if not sorted_values:
        return 0.0
    rank = (len(sorted_values) - 1) * percentile
    lower = int(rank)
    upper = min(lower + 1, len(sorted_values) - 1)
    weight = rank - lower
    return sorted_values[lower] * (1 - weight) + sorted_values[upper] * weight


def _base_request(
    *, payload: dict[str, Any], policy_context: dict[str, Any]
) -> dict[str, Any]:
    return {
        "workspace_id": "workspace-demo",
        "principal": {
            "type": "token",
            "id": "api:perf-bot",
        },
        "object_type": "document",
        "operation": "create",
        "payload": payload,
        "policy_context": policy_context,
    }


def _scenario_governed_no_matching_policy() -> dict[str, Any]:
    return _base_request(
        payload={"title": "Q4 Plan"},
        policy_context={
            "governed": True,
            "selected_policies": [],
        },
    )


def _scenario_governed_missing_field_denied() -> dict[str, Any]:
    return _base_request(
        payload={},
        policy_context={
            "governed": True,
            "selected_policies": [
                {
                    "policy_key": "document_create_default",
                    "policy_version": 1,
                    "effect": "allow",
                    "rules": [
                        {
                            "rule_key": "require_title",
                            "rule_definition": {
                                "type": "required_fields",
                                "fields": ["title"],
                            },
                            "reason_code": "missing_required_fields",
                            "reason_message": "title is required.",
                        }
                    ],
                }
            ],
        },
    )


def _scenario_governed_approved() -> dict[str, Any]:
    return _base_request(
        payload={"title": "Q4 Plan"},
        policy_context={
            "governed": True,
            "selected_policies": [
                {
                    "policy_key": "document_create_default",
                    "policy_version": 1,
                    "effect": "allow",
                    "rules": [
                        {
                            "rule_key": "require_title",
                            "rule_definition": {
                                "type": "required_fields",
                                "fields": ["title"],
                            },
                            "reason_code": "missing_required_fields",
                            "reason_message": "title is required.",
                        }
                    ],
                }
            ],
        },
    )


def _scenario_out_of_scope_approved() -> dict[str, Any]:
    return _base_request(
        payload={"name": "Acme"},
        policy_context={
            "governed": False,
            "selected_policies": [],
        },
    )


SCENARIO_BUILDERS = {
    "governed_no_matching_policy": _scenario_governed_no_matching_policy,
    "governed_missing_field_denied": _scenario_governed_missing_field_denied,
    "governed_approved": _scenario_governed_approved,
    "out_of_scope_approved": _scenario_out_of_scope_approved,
}


def _run_scenario(
    name: str, request: dict[str, Any], *, iterations: int, warmup: int
) -> dict[str, Any]:
    for index in range(warmup):
        evaluate_execute(request, idempotency_key=f"{name}-warmup-{index}")

    latencies_ms: list[float] = []
    started_at = time.perf_counter()

    for index in range(iterations):
        tick = time.perf_counter_ns()
        evaluate_execute(request, idempotency_key=f"{name}-{index}")
        elapsed_ns = time.perf_counter_ns() - tick
        latencies_ms.append(elapsed_ns / 1_000_000)

    duration_sec = time.perf_counter() - started_at
    sorted_latencies = sorted(latencies_ms)

    return {
        "scenario": name,
        "iterations": iterations,
        "warmup": warmup,
        "duration_sec": round(duration_sec, 6),
        "throughput_rps": round(iterations / duration_sec, 2)
        if duration_sec > 0
        else 0.0,
        "latency_ms": {
            "p50": round(_percentile(sorted_latencies, 0.50), 6),
            "p95": round(_percentile(sorted_latencies, 0.95), 6),
            "p99": round(_percentile(sorted_latencies, 0.99), 6),
            "max": round(sorted_latencies[-1], 6) if sorted_latencies else 0.0,
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Benchmark Kernite execute evaluation latency/throughput."
    )
    parser.add_argument(
        "--iterations",
        type=int,
        default=20000,
        help="Number of benchmark iterations per scenario.",
    )
    parser.add_argument(
        "--warmup",
        type=int,
        default=500,
        help="Warmup iterations per scenario.",
    )
    parser.add_argument(
        "--scenario",
        choices=["all", *SCENARIO_BUILDERS.keys()],
        default="all",
        help="Scenario to run.",
    )

    args = parser.parse_args()

    scenario_names = (
        list(SCENARIO_BUILDERS.keys()) if args.scenario == "all" else [args.scenario]
    )

    report = {
        "benchmark": "kernite.execute",
        "timestamp_unix": int(time.time()),
        "python": {
            "version": platform.python_version(),
        },
        "results": [],
    }

    for scenario_name in scenario_names:
        request = SCENARIO_BUILDERS[scenario_name]()
        report["results"].append(
            _run_scenario(
                scenario_name,
                request,
                iterations=args.iterations,
                warmup=args.warmup,
            )
        )

    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
