from __future__ import annotations

import json
from collections.abc import Callable

from env.episodes import HAND_AUTHORED_SCENARIOS, get_hand_authored_scenario
from env.models import SchedulingAction
from env.simulator import GPUSchedulingSimulator
from training.heuristic_baselines import POLICIES


PolicyFn = Callable[[GPUSchedulingSimulator], SchedulingAction]


def run_policy(policy_name: str, policy: PolicyFn, scenario_id: str) -> dict[str, float | int | str]:
    simulator = GPUSchedulingSimulator(get_hand_authored_scenario(scenario_id))
    done = False
    while not done:
        result = simulator.step(policy(simulator))
        done = result.done
    metrics = simulator.metrics
    return {
        "policy": policy_name,
        "scenario_id": scenario_id,
        "total_reward": round(metrics.total_reward, 2),
        "completed_jobs": metrics.completed_jobs,
        "missed_deadlines": metrics.missed_deadlines,
        "invalid_actions": metrics.invalid_actions,
        "allocation": round(metrics.allocation, 3),
        "total_cost": round(metrics.total_cost, 2),
        "average_pending_ticks": round(metrics.average_pending_ticks, 2),
    }


def generate_report() -> list[dict[str, float | int | str]]:
    rows: list[dict[str, float | int | str]] = []
    for scenario_id in HAND_AUTHORED_SCENARIOS:
        for policy_name, policy in POLICIES.items():
            rows.append(run_policy(policy_name, policy, scenario_id))
    return rows


def format_markdown_table(rows: list[dict[str, float | int | str]]) -> str:
    headers = list(rows[0].keys())
    table = ["| " + " | ".join(headers) + " |", "| " + " | ".join(["---"] * len(headers)) + " |"]
    for row in rows:
        table.append("| " + " | ".join(str(row[h]) for h in headers) + " |")
    return "\n".join(table)


if __name__ == "__main__":
    report = generate_report()
    print(format_markdown_table(report))
    print()
    print(json.dumps(report, indent=2))