from __future__ import annotations

from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any

from app.modules.agent.engine.orchestrator.artifact_store import ArtifactStore
from app.modules.agent.engine.orchestrator.evidence_store import EvidenceStore
from app.modules.agent.engine.orchestrator.models import ExecutionPlan, TaskSpec

ProgressCallback = Callable[[str, str, str, dict | None], Awaitable[None] | None]
GraphResolver = Callable[[str, str], Any]
GraphInvoker = Callable[[Any, dict, str], dict]


@dataclass
class ExecutionContext:
    task: TaskSpec
    plan: ExecutionPlan
    graph_resolver: GraphResolver
    graph_invoker: GraphInvoker
    progress_cb: ProgressCallback | None = None
    artifacts: ArtifactStore | None = None
    evidences: EvidenceStore | None = None

    def __post_init__(self) -> None:
        if self.artifacts is None:
            self.artifacts = ArtifactStore()
        if self.evidences is None:
            self.evidences = EvidenceStore()
