from __future__ import annotations

from typing import Generic, Sequence, TypeVar

from app.core.agent.utils.workflow.context import WorkflowContext
from app.core.agent.utils.workflow.step import WorkflowStep


TContext = TypeVar("TContext", bound=WorkflowContext)


class WorkflowGraph(Generic[TContext]):
    def __init__(self, workflow_id: str, source: str, steps: Sequence[WorkflowStep[TContext]]) -> None:
        self._workflow_id = workflow_id
        self._source = source
        self._steps = tuple(steps)

    async def run(self, context: TContext) -> TContext:
        trace = context.runtime.trace.module(self._source)
        trace.log("workflow_started", {"workflow_id": self._workflow_id})
        for step in self._steps:
            context = await self._run_step(context, step)
        trace.log("workflow_completed", {"workflow_id": self._workflow_id})
        return context

    async def _run_step(self, context: TContext, step: WorkflowStep[TContext]) -> TContext:
        request_id = context.runtime.request.request_id
        trace = context.runtime.trace.module(self._source)
        trace.log(
            "step_started",
            {"workflow_id": self._workflow_id, "step_id": step.step_id, "input": step.trace_input(context)},
        )
        await context.runtime.publisher.publish_status(
            request_id,
            self._source,
            f"Шаг workflow: {step.title}.",
            {"workflow_id": self._workflow_id, "step_id": step.step_id},
        )
        context = await step.run(context)
        trace.log(
            "step_completed",
            {"workflow_id": self._workflow_id, "step_id": step.step_id, "output": step.trace_output(context)},
        )
        return context
