import argparse
import csv
import json
import time
from pathlib import Path

from concurrency_test import DEFAULT_NEW_ENDPOINT, _encode_image, run_endpoint


def parse_levels(value: str) -> list[int]:
    return [int(item.strip()) for item in value.split(",") if item.strip()]


def write_csv(path: Path, rows: list[dict]) -> None:
    columns = [
        "concurrency",
        "requests",
        "success",
        "failed",
        "wall_seconds",
        "throughput_rps",
        "latency_min",
        "latency_mean",
        "latency_p50",
        "latency_p95",
        "latency_max",
    ]
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=columns)
        writer.writeheader()
        for row in rows:
            lat = row["latency_seconds"]
            writer.writerow(
                {
                    "concurrency": row["concurrency"],
                    "requests": row["requests"],
                    "success": row["success"],
                    "failed": row["failed"],
                    "wall_seconds": row["wall_seconds"],
                    "throughput_rps": row["throughput_rps"],
                    "latency_min": lat["min"],
                    "latency_mean": lat["mean"],
                    "latency_p50": lat["p50"],
                    "latency_p95": lat["p95"],
                    "latency_max": lat["max"],
                }
            )


def write_plot(path: Path, rows: list[dict]) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    x = [row["concurrency"] for row in rows]
    mean = [row["latency_seconds"]["mean"] for row in rows]
    p50 = [row["latency_seconds"]["p50"] for row in rows]
    p95 = [row["latency_seconds"]["p95"] for row in rows]

    plt.figure(figsize=(10, 6))
    plt.plot(x, mean, marker="o", label="mean")
    plt.plot(x, p50, marker="o", label="p50")
    plt.plot(x, p95, marker="o", label="p95")
    plt.xlabel("Concurrency")
    plt.ylabel("Latency per request (seconds)")
    plt.title("vLLM OCR latency vs concurrency")
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(path, dpi=160)


def main() -> None:
    parser = argparse.ArgumentParser(description="Sweep vLLM OCR endpoint concurrency levels.")
    parser.add_argument("--image", required=True)
    parser.add_argument("--url", default=DEFAULT_NEW_ENDPOINT)
    parser.add_argument("--levels", default="1,2,4,6,8,10,12,20,30,40,50")
    parser.add_argument("--timeout", type=float, default=900)
    parser.add_argument("--output-json", default="concurrency_sweep_results.json")
    parser.add_argument("--output-csv", default="concurrency_sweep_results.csv")
    parser.add_argument("--output-plot", default="concurrency_sweep_latency.png")
    args = parser.parse_args()

    image_path = Path(args.image)
    image_type = image_path.suffix.lstrip(".").lower() or "png"
    payload = {
        "file": _encode_image(image_path),
        "type": "jpg" if image_type == "jpeg" else image_type,
        "skip_text_detection": False,
        "skip_table_detection": False,
        "recognize_math": False,
        "ocr_with_boxes": True,
    }

    levels = parse_levels(args.levels)
    rows = []
    for concurrency in levels:
        print(f"running concurrency={concurrency}", flush=True)
        rows.append(
            run_endpoint(
                "vllm",
                args.url,
                payload,
                requests_count=concurrency,
                concurrency=concurrency,
                timeout=args.timeout,
            )
        )

    result = {
        "started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "image": str(image_path),
        "url": args.url,
        "levels": levels,
        "results": rows,
    }
    Path(args.output_json).write_text(json.dumps(result, indent=2), encoding="utf-8")
    write_csv(Path(args.output_csv), rows)
    write_plot(Path(args.output_plot), rows)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()