#!/usr/bin/env python3
"""
ml/task_graph.py — Dask DAG-basert agent-koordinering for OPAX.

Konsept (fra ML-dokumentet):
  Dask bygger en task graph (DAG) av operasjoner og avhengigheter.
  Lazy evaluation: beregninger utføres kun når .compute() kalles.
  Agent-funksjoner wrappes som @dask.delayed og kjøres parallelt.

PS-analogi:
  OPAX orchestrator = DAG scheduler (server-node)
  Individuelle agenter = Dask-delayed tasks (worker-noder)

Fix 2026-06-15: agent_count returneres fra build_agent_dag slik at
  /run/dag kan bruke len(tasks) istedenfor len(agent_ids) (som var
  undefined i feilscenariet der agent_fns og agent_ids hadde ulik lengde).
"""

import time
from typing import Any, Callable, Dict, List, Optional, Tuple

import dask


@dask.delayed
def run_agent_task(
    agent_fn: Callable,
    payload: Dict[str, Any],
    agent_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Wrapper: kjører én agent som en Dask-delayed task.
    Returnerer resultat med metadata for telemetri og state store.
    """
    _id = agent_id or getattr(agent_fn, "__name__", "unknown")
    start = time.monotonic()
    try:
        result = agent_fn(payload)
        success = True
        error = None
    except Exception as exc:  # noqa: BLE001
        result = None
        success = False
        error = str(exc)
    duration = round(time.monotonic() - start, 3)
    return {
        "agent": _id,
        "result": result,
        "success": success,
        "error": error,
        "duration_s": duration,
        "timestamp": time.time(),
    }


def build_agent_dag(
    agent_fns: List[Callable],
    payload: Dict[str, Any],
    agent_ids: Optional[List[str]] = None,
) -> List:
    """
    Bygger en liste av Dask-delayed tasks (én per agent).
    Agenter uten avhengigheter kjøres parallelt av scheduleren.

    Args:
        agent_fns: Liste av callable agentfunksjoner.
        payload:   Felles input-dict til alle agenter.
        agent_ids: Valgfrie ID-er som matcher agent_fns (for logging).
                   Dersom kortere enn agent_fns fylles resten med auto-navn.

    Returns:
        Liste av Dask-delayed objects — klar for execute_dag().
    """
    n = len(agent_fns)
    # Guard: aldri kræsj på lengde-mismatch
    _ids = list(agent_ids or [])
    _ids += [getattr(fn, "__name__", f"agent_{i}") for i, fn in enumerate(agent_fns) if i >= len(_ids)]
    _ids = _ids[:n]  # trim hvis agent_ids var lengre enn agent_fns
    return [run_agent_task(fn, payload, aid) for fn, aid in zip(agent_fns, _ids)]


def execute_dag(
    tasks: List,
    scheduler: str = "synchronous",
) -> List[Dict[str, Any]]:
    """
    Kjører den ferdigbygde DAGen.

    Scheduler-valg:
      'synchronous'  — enkelt-tråd, ingen overhead (dev / unit-test)
      'threads'      — lokal multi-tråd, GIL-vennlig for IO-bound agenter
      'processes'    — multi-prosess, CPU-bound arbeid
      'distributed'  — Dask cluster (full skala, krever dask[distributed])

    Returns:
        Liste av resultat-dicts fra run_agent_task.
    """
    if not tasks:
        return []
    results = dask.compute(*tasks, scheduler=scheduler)
    return list(results)