"""Workload protocol and registry."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable, Protocol, runtime_checkable

import numpy as np

from benchmarks.adapters.base import DatabaseAdapter
from benchmarks.config import ScalePreset


@runtime_checkable
class Workload(Protocol):
    """Protocol for a single benchmark workload."""

    @property
    def workload_id(self) -> str:
        """Unique ID like 'W01_node_creation'."""
        ...

    @property
    def description(self) -> str:
        """Human-readable description."""
        ...

    def setup(self, adapter: DatabaseAdapter, scale: ScalePreset) -> None:
        """Per-workload setup (called once before warm-up)."""
        ...

    def run(self, adapter: DatabaseAdapter) -> Any:
        """Execute the workload once. Return value is recorded for correctness."""
        ...

    def cleanup(self, adapter: DatabaseAdapter) -> None:
        """Per-workload cleanup (called once after all iterations)."""
        ...


@dataclass
class WorkloadSpec:
    """Concrete workload specification with parameters."""

    workload_id: str
    description: str
    setup_fn: Callable[[DatabaseAdapter, ScalePreset], None] | None = None
    run_fn: Callable[[DatabaseAdapter], Any] | None = None
    cleanup_fn: Callable[[DatabaseAdapter], None] | None = None
    params: dict[str, Any] = field(default_factory=dict)


# Global registry populated by workload modules
WORKLOAD_REGISTRY: dict[str, WorkloadSpec] = {}


def register_workload(spec: WorkloadSpec) -> WorkloadSpec:
    """Register a workload spec in the global registry."""
    WORKLOAD_REGISTRY[spec.workload_id] = spec
    return spec
