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 FindFilesWorkflowContext
from app.core.agent.processes.v2.workflows.doc_explain_find_files.workflow_runtime.pipeline_logging import log_pipeline_step, log_ranking
from app.core.agent.utils.process_v2.evidence.assembler import DocsEvidenceAssembler
from app.core.agent.utils.workflow import WorkflowStep

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


class BuildFilesEvidenceStep(WorkflowStep[TContext], Generic[TContext]):
    step_id = "build_files_evidence"
    title = "Сборка file evidence"

    def __init__(self, assembler: DocsEvidenceAssembler) -> None:
        self._assembler = assembler

    async def run(self, context: TContext) -> TContext:
        if context.answer:
            return context
        context.files = self._assembler.assemble_files(context.rows, context.route)
        context.runtime.trace.module("process.v2.evidence").log(
            "evidence_assembled",
            {
                "mode": "find_files",
                "file_count": len(context.files),
                "files": [item.path for item in context.files],
            },
        )
        log_pipeline_step(
            context.runtime,
            "evidence_assembled",
            {
                "mode": "find_files",
                "primary_file": context.files[0].path if context.files else None,
                "file_count": len(context.files),
            },
        )
        log_ranking(context.runtime, context.files)
        return context

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