"""Section 17-19: compare agent configurations across the whole case suite."""

import asyncio
import json
import os
import statistics
import sys
from datetime import datetime

from temporalio.client import Client

from cases import CASES, LABELS
from workflow import BenchmarkWorkflow

TRIAGE_PROMPT = (
    "You are a server hardware failure-triage expert. Reply with exactly one "
    "label and nothing else."
)

CONFIGS = [
    {
        "name": "A_full_context_256",
        "max_context_chars": 5000,
        "max_tokens": 256,
        "temperature": 0,
        "model": "sglang-rtx6000",
    },
    {
        "name": "B_reduced_context_256",
        "max_context_chars": 250,
        "max_tokens": 256,
        "temperature": 0,
        "model": "sglang-rtx6000",
    },
    {
        "name": "C_reduced_context_64",
        "max_context_chars": 250,
        "max_tokens": 64,
        "temperature": 0,
        "model": "sglang-rtx6000",
    },
    {
        "name": "D_specialized_prompt_64_nothink",
        "max_context_chars": 250,
        "max_tokens": 64,
        "temperature": 0,
        "system_prompt": TRIAGE_PROMPT,
        "enable_thinking": False,
        "model": "sglang-rtx6000",
    },
    {
        "name": "E_alt_model_3090_nothink",
        "max_context_chars": 250,
        "max_tokens": 64,
        "temperature": 0,
        "system_prompt": TRIAGE_PROMPT,
        "enable_thinking": False,
        "model": "llamacpp-3090",
    },
]

SUITE = {
    "name": "server-failure-triage-v1",
    "label_set": LABELS,
    "cases": CASES,
    "configs": CONFIGS,
}


def summarize(rows: list) -> list:
    by_config: dict = {}
    for row in rows:
        by_config.setdefault(row["name"], []).append(row)

    summary = []
    for name, group in by_config.items():
        latencies = sorted(r["latency_seconds"] for r in group)
        passes = [r for r in group if r["passed"]]
        total_tokens = sum(r["total_tokens"] or 0 for r in group)
        summary.append(
            {
                "config": name,
                "backend": group[0].get("backend"),
                "cases": len(group),
                "accuracy": len(passes) / len(group),
                "mean_latency_s": statistics.fmean(latencies),
                "median_latency_s": statistics.median(latencies),
                "p95_latency_s": latencies[min(len(latencies) - 1, int(0.95 * len(latencies)))],
                "mean_prompt_tokens": statistics.fmean(r["prompt_tokens"] or 0 for r in group),
                "mean_completion_tokens": statistics.fmean(r["completion_tokens"] or 0 for r in group),
                "mean_reasoning_tokens": statistics.fmean(r.get("reasoning_tokens") or 0 for r in group),
                "total_tokens": total_tokens,
                "tokens_per_correct_answer": (total_tokens / len(passes)) if passes else None,
                "truncated": sum(1 for r in group if r.get("finish_reason") == "length"),
            }
        )
    return summary


def to_markdown(summary: list) -> str:
    head = (
        "| Configuration | Accuracy | Median latency | P95 latency | "
        "Prompt tok | Completion tok | Tokens / correct |\n"
        "|---|---:|---:|---:|---:|---:|---:|\n"
    )
    rows = []
    for s in summary:
        tpc = f"{s['tokens_per_correct_answer']:.0f}" if s["tokens_per_correct_answer"] else "n/a"
        rows.append(
            f"| {s['config']} | {s['accuracy'] * 100:.1f}% | {s['median_latency_s']:.2f} s | "
            f"{s['p95_latency_s']:.2f} s | {s['mean_prompt_tokens']:.0f} | "
            f"{s['mean_completion_tokens']:.0f} | {tpc} |"
        )
    return head + "\n".join(rows) + "\n"


async def main():
    workflow_id = sys.argv[1] if len(sys.argv) > 1 else "benchmark-triage-001"
    client = await Client.connect("localhost:7233")

    handle = await client.start_workflow(
        BenchmarkWorkflow.run,
        SUITE,
        id=workflow_id,
        task_queue="agent-optimization",
    )
    print(f"Started {workflow_id} ({len(CASES)} cases x {len(CONFIGS)} configs)")

    result = await handle.result()
    summary = summarize(result["rows"])

    os.makedirs("results", exist_ok=True)
    payload = {
        "suite": result["suite"],
        "generated_at": datetime.now().isoformat(timespec="seconds"),
        "workflow_id": workflow_id,
        "summary": summary,
        "rows": result["rows"],
    }
    with open(f"results/{workflow_id}.json", "w") as fh:
        json.dump(payload, fh, indent=2)

    table = to_markdown(summary)
    with open(f"results/{workflow_id}.md", "w") as fh:
        fh.write(table)

    print(table)
    print(f"Saved results/{workflow_id}.json and results/{workflow_id}.md")


if __name__ == "__main__":
    asyncio.run(main())