from __future__ import annotations

from typing import Generic, TypeVar

from app.core.agent.processes.v2.workflows.general_qa_summary.steps.retrieval.candidate_rows import CandidateRowsBuilder
from app.core.agent.processes.v2.workflows.general_qa_summary.workflow_runtime.context_protocols import RetrievalWorkflowContext
from app.core.agent.processes.v2.workflows.general_qa_summary.workflow_runtime.pipeline_logging import log_retrieval_trace
from app.core.agent.utils.workflow import WorkflowStep

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


class PrepareCandidateRowsStep(WorkflowStep[TContext], Generic[TContext]):
    step_id = "prepare_candidate_rows"
    title = "Подготовка candidate rows"

    def __init__(self, builder: CandidateRowsBuilder) -> None:
        self._builder = builder

    async def run(self, context: TContext) -> TContext:
        if context.answer or context.retrieval_plan is None:
            return context
        prepared = self._builder.build(context.retrieved_rows, context.route)
        context.rows = prepared.rows
        log_retrieval_trace(
            context.runtime,
            context.route,
            context.retrieval_plan,
            context.retrieved_rows,
            prepared.metadata_rows,
            prepared.rows,
        )
        return context

    def trace_output(self, context: TContext) -> dict[str, object]:
        return {"row_count": len(context.rows)}
