from __future__ import annotations

from collections.abc import Callable
from enum import Enum
from typing import Any

from kernbench.common.types import Completion, SimEngine, Trace

from .context import RuntimeContext
from .types import BenchResult, DeviceSelector


class CompletionPolicy(str, Enum):
    LAST_SUBMITTED = "last_submitted"
    LAST_COMPLETED = "last_completed"  # requires trace/timestamps or engine support; stub for now
    ALL_OK_FAIL_FAST = "all_ok_fail_fast"


BenchFn = Callable[[RuntimeContext], Any]
EngineFactory = Callable[[object, DeviceSelector], SimEngine]


def run_bench(
    *,
    topology: object,
    bench_fn: BenchFn,
    device: DeviceSelector,
    engine_factory: EngineFactory,
    correlation_id: str = "bench0",
    completion_policy: CompletionPolicy = CompletionPolicy.LAST_SUBMITTED,
) -> BenchResult:
    """Minimal bench runner.

    - topology: compiled topology object (opaque to runtime here)
    - bench_fn: callable ``run(torch)`` receiving a RuntimeContext
    - device: DeviceSelector ("all" or "sip:<N>")
    - engine_factory: builds sim_engine for given topology & device
    - completion_policy: how to determine overall completion/result
    """
    engine = engine_factory(topology, device)
    # Extract spec from TopologyHandle or TopologyGraph
    topo_obj = getattr(topology, "topology_obj", topology)
    spec = getattr(topo_obj, "spec", None)
    ctx = RuntimeContext(
        engine=engine, target_device=device,
        correlation_id=correlation_id, spec=spec,
    )

    bench_fn(ctx)
    ctx.wait_all()

    collected_traces = ctx._traces or None

    handles = ctx.handles()
    if not handles:
        return BenchResult(
            completion=Completion(
                ok=False, error_code="NO_REQUESTS", error_message="Bench submitted no requests"
            ),
            correlation_id=correlation_id,
            trace=None,
            traces=collected_traces,
            engine=engine,
        )

    if completion_policy == CompletionPolicy.LAST_SUBMITTED:
        last = handles[-1]
        completion, trace = engine.get_completion(last)
        return BenchResult(
            completion=completion, correlation_id=correlation_id,
            trace=trace, traces=collected_traces, engine=engine,
        )

    if completion_policy == CompletionPolicy.ALL_OK_FAIL_FAST:
        last_trace: Trace | None = None
        for h in handles:
            c, t = engine.get_completion(h)
            last_trace = t if t is not None else last_trace
            if not c.ok:
                return BenchResult(
                    completion=c, correlation_id=correlation_id,
                    trace=last_trace, traces=collected_traces, engine=engine,
                )
        return BenchResult(
            completion=Completion(ok=True), correlation_id=correlation_id,
            trace=last_trace, traces=collected_traces, engine=engine,
        )

    # LAST_COMPLETED placeholder (needs engine support for timing). Fall back.
    last = handles[-1]
    completion, trace = engine.get_completion(last)
    return BenchResult(
        completion=completion, correlation_id=correlation_id,
        trace=trace, traces=collected_traces, engine=engine,
    )