#!/usr/bin/env python3
"""Helpers for adaptive-concurrency serving benchmarks."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path
from typing import Any, Iterable


POINT_FIELDS = [
    "timestamp",
    "engine",
    "tp",
    "dp",
    "mark",
    "isl",
    "osl",
    "concurrency",
    "num_prompts",
    "warmup_requests",
    "attempt",
    "status",
    "completed",
    "failed",
    "duration_s",
    "request_tps",
    "input_tps",
    "output_tps",
    "total_tps",
    "mean_input_tokens",
    "mean_output_tokens",
    "ttft_p50_ms",
    "ttft_p95_ms",
    "ttft_p99_ms",
    "tpot_p50_ms",
    "tpot_p95_ms",
    "tpot_p99_ms",
    "e2e_p50_ms",
    "e2e_p95_ms",
    "e2e_p99_ms",
    "itl_p50_ms",
    "itl_p95_ms",
    "itl_p99_ms",
    "gain_pct",
    "plateau_streak",
    "error_type",
    "validation_errors",
    "raw_file",
    "detail_log",
]


SHAPE_FIELDS = [
    "timestamp",
    "engine",
    "tp",
    "dp",
    "mark",
    "isl",
    "osl",
    "status",
    "stop_reason",
    "tested_points",
    "search_cap",
    "max_successful_concurrency",
    "saturation_concurrency",
    "stop_probe_concurrency",
    "best_tps_concurrency",
    "best_total_tps",
    "last_total_tps",
]


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    if not path.exists():
        return rows
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                value = json.loads(line)
            except json.JSONDecodeError:
                continue
            if isinstance(value, dict):
                rows.append(value)
    return rows


def last_json_object(path: Path) -> dict[str, Any]:
    rows = read_jsonl(path)
    if not rows:
        raise ValueError(f"no valid JSON object in {path}")
    return rows[-1]


def percentile_value(data: dict[str, Any], prefix: str, percentile: str) -> float:
    if percentile == "p50":
        key = f"median_{prefix}_ms"
    else:
        key = f"{percentile}_{prefix}_ms"
    return float(data.get(key, 0.0) or 0.0)


def command_shapes(args: argparse.Namespace) -> int:
    with args.matrix.open("r", encoding="utf-8") as f:
        data = json.load(f)

    mode = args.mode or data.get("mode", "Y")
    print("mark\tinput_len\toutput_len")
    for isl_text in sorted(data["matrix"], key=int):
        osl_map = data["matrix"][isl_text]
        for osl_text in sorted(osl_map, key=int):
            mark = osl_map[osl_text]
            if mode == "Y" and mark != "Y":
                continue
            if mode == "Y+P" and mark not in ("Y", "P"):
                continue
            if mode != "all" and mark == "N":
                continue
            print(f"{mark}\t{int(isl_text)}\t{int(osl_text)}")
    return 0


def command_dataset_capacity(args: argparse.Namespace) -> int:
    try:
        with args.path.open("r", encoding="utf-8") as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError) as exc:
        raise SystemExit(f"invalid dataset {args.path}: {exc}")
    if not isinstance(data, list):
        raise SystemExit(f"dataset must be a JSON list: {args.path}")
    capacity = 0
    for row in data:
        if not isinstance(row, dict):
            continue
        conversations = row.get("conversations", row.get("conversation", []))
        if isinstance(conversations, list) and len(conversations) >= 2:
            capacity += 1
    print(capacity)
    return 0


def command_parse_result(args: argparse.Namespace) -> int:
    validation_errors: list[str] = []
    try:
        data = last_json_object(args.input)
    except (OSError, ValueError) as exc:
        data = {}
        validation_errors.append(str(exc))

    completed = int(data.get("completed", 0) or 0)
    input_lens = [int(v) for v in data.get("input_lens", []) if v is not None]
    output_lens = [int(v) for v in data.get("output_lens", []) if v is not None]
    errors = [str(v) for v in data.get("errors", []) if str(v)]
    mean_input = sum(input_lens) / len(input_lens) if input_lens else 0.0
    mean_output = sum(output_lens) / len(output_lens) if output_lens else 0.0

    if completed != args.expected_prompts:
        validation_errors.append(
            f"completed={completed}, expected={args.expected_prompts}"
        )
    if errors:
        validation_errors.append(f"request_errors={len(errors)}")

    input_low = args.isl * (1.0 - args.input_tolerance_pct / 100.0)
    input_high = args.isl * (1.0 + args.input_tolerance_pct / 100.0)
    if not input_low <= mean_input <= input_high:
        validation_errors.append(
            f"mean_input_tokens={mean_input:.2f}, expected_range={input_low:.2f}-{input_high:.2f}"
        )

    output_low = args.osl * (1.0 - args.output_tolerance_pct / 100.0)
    output_high = args.osl * (1.0 + args.output_tolerance_pct / 100.0)
    if not output_low <= mean_output <= output_high:
        validation_errors.append(
            f"mean_output_tokens={mean_output:.2f}, expected_range={output_low:.2f}-{output_high:.2f}"
        )

    status = "COMPLETED" if not validation_errors else "INVALID_WORKLOAD"
    result = {
        "status": status,
        "completed": completed,
        "failed": max(args.expected_prompts - completed, len(errors)),
        "duration_s": float(data.get("duration", 0.0) or 0.0),
        "request_tps": float(data.get("request_throughput", 0.0) or 0.0),
        "input_tps": float(data.get("input_throughput", 0.0) or 0.0),
        "output_tps": float(data.get("output_throughput", 0.0) or 0.0),
        "total_tps": float(data.get("total_throughput", 0.0) or 0.0),
        "mean_input_tokens": mean_input,
        "mean_output_tokens": mean_output,
        "ttft_p50_ms": percentile_value(data, "ttft", "p50"),
        "ttft_p95_ms": percentile_value(data, "ttft", "p95"),
        "ttft_p99_ms": percentile_value(data, "ttft", "p99"),
        "tpot_p50_ms": percentile_value(data, "tpot", "p50"),
        "tpot_p95_ms": percentile_value(data, "tpot", "p95"),
        "tpot_p99_ms": percentile_value(data, "tpot", "p99"),
        "e2e_p50_ms": percentile_value(data, "e2e_latency", "p50"),
        "e2e_p95_ms": percentile_value(data, "e2e_latency", "p95"),
        "e2e_p99_ms": percentile_value(data, "e2e_latency", "p99"),
        "itl_p50_ms": percentile_value(data, "itl", "p50"),
        "itl_p95_ms": percentile_value(data, "itl", "p95"),
        "itl_p99_ms": percentile_value(data, "itl", "p99"),
        "validation_errors": validation_errors,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)
        f.write("\n")
    print(status)
    return 0 if status == "COMPLETED" else 2


def command_gain(args: argparse.Namespace) -> int:
    if args.previous <= 0:
        gain_pct = float("inf")
        meaningful = True
    else:
        gain_pct = (args.current - args.previous) / args.previous * 100.0
        meaningful = gain_pct >= args.threshold_pct
    gain_text = "inf" if gain_pct == float("inf") else f"{gain_pct:.6f}"
    print(f"{gain_text}\t{1 if meaningful else 0}")
    return 0


def write_csv(path: Path, rows: Iterable[dict[str, Any]], fields: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            output = dict(row)
            if isinstance(output.get("validation_errors"), list):
                output["validation_errors"] = "; ".join(output["validation_errors"])
            writer.writerow(output)


def command_summarize(args: argparse.Namespace) -> int:
    points = read_jsonl(args.points)
    shapes = read_jsonl(args.shapes)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    write_csv(args.output_dir / "adaptive_points.csv", points, POINT_FIELDS)
    write_csv(args.output_dir / "adaptive_summary.csv", shapes, SHAPE_FIELDS)

    summary_jsonl = args.output_dir / "adaptive_summary.jsonl"
    with summary_jsonl.open("w", encoding="utf-8") as f:
        for row in shapes:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    report = args.output_dir / "adaptive_summary.md"
    with report.open("w", encoding="utf-8") as f:
        f.write("# Adaptive concurrency search summary\n\n")
        f.write(
            "| Engine | TP | DP | ISL | OSL | Stop | Saturation C | Best TPS C | Best Total TPS | Max successful C |\n"
        )
        f.write("|---|---:|---:|---:|---:|---|---:|---:|---:|---:|\n")
        for row in shapes:
            f.write(
                "| {engine} | {tp} | {dp} | {isl} | {osl} | {stop_reason} | "
                "{saturation_concurrency} | {best_tps_concurrency} | {best_total_tps} | "
                "{max_successful_concurrency} |\n".format(**{k: row.get(k, "") for k in SHAPE_FIELDS})
            )
        f.write("\n")
        f.write(
            "`Saturation C` is the first point in the final low-gain streak. "
            "`Best TPS C` is the tested point with the highest observed Total TPS.\n"
        )
    print(f"wrote summaries under {args.output_dir}")
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    shapes = subparsers.add_parser("shapes")
    shapes.add_argument("--matrix", type=Path, required=True)
    shapes.add_argument("--mode", choices=["Y", "Y+P", "all"], default=None)
    shapes.set_defaults(func=command_shapes)

    dataset_capacity = subparsers.add_parser("dataset-capacity")
    dataset_capacity.add_argument("--path", type=Path, required=True)
    dataset_capacity.set_defaults(func=command_dataset_capacity)

    parse_result = subparsers.add_parser("parse-result")
    parse_result.add_argument("--input", type=Path, required=True)
    parse_result.add_argument("--output", type=Path, required=True)
    parse_result.add_argument("--expected-prompts", type=int, required=True)
    parse_result.add_argument("--isl", type=int, required=True)
    parse_result.add_argument("--osl", type=int, required=True)
    parse_result.add_argument("--input-tolerance-pct", type=float, default=5.0)
    parse_result.add_argument("--output-tolerance-pct", type=float, default=10.0)
    parse_result.set_defaults(func=command_parse_result)

    gain = subparsers.add_parser("gain")
    gain.add_argument("--previous", type=float, required=True)
    gain.add_argument("--current", type=float, required=True)
    gain.add_argument("--threshold-pct", type=float, required=True)
    gain.set_defaults(func=command_gain)

    summarize = subparsers.add_parser("summarize")
    summarize.add_argument("--points", type=Path, required=True)
    summarize.add_argument("--shapes", type=Path, required=True)
    summarize.add_argument("--output-dir", type=Path, required=True)
    summarize.set_defaults(func=command_summarize)
    return parser


def main() -> int:
    args = build_parser().parse_args()
    return int(args.func(args))


if __name__ == "__main__":
    raise SystemExit(main())
