from __future__ import annotations

from typing import Generic, TypeVar

from app.core.agent.processes.v2.workflows.doc_explain_find_files.workflow_runtime.context_protocols import RetrievalWorkflowContext
from app.core.agent.processes.v2.workflows.doc_explain_find_files.workflow_runtime.pipeline_logging import log_pipeline_step
from app.core.agent.utils.process_v2.plan_resolver import RetrievalPlanResolver
from app.core.agent.utils.workflow import WorkflowStep

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


class ResolveRetrievalPlanStep(WorkflowStep[TContext], Generic[TContext]):
    step_id = "resolve_retrieval_plan"
    title = "Выбор retrieval-плана"

    def __init__(self, resolver: RetrievalPlanResolver) -> None:
        self._resolver = resolver

    async def run(self, context: TContext) -> TContext:
        if context.answer:
            return context
        plan = self._resolver.resolve(context.route)
        context.retrieval_plan = plan
        context.runtime.trace.module("process.v2.retrieval_policy").log(
            "retrieval_plan_resolved",
            {"profile": plan.profile, "layers": plan.layers, "limit": plan.limit, "filters": plan.filters},
        )
        log_pipeline_step(
            context.runtime,
            "retrieval_profile_selected",
            {"profile": plan.profile, "layers": plan.layers, "filters": plan.filters},
        )
        return context

    def trace_output(self, context: TContext) -> dict[str, object]:
        return {"profile": getattr(context.retrieval_plan, "profile", "")}
