from __future__ import annotations

import asyncio

from app.core.agent.processes.v2.workflows.general_qa_summary.workflow_runtime.context import GeneralQaSummaryContext
from app.core.agent.utils.llm import AgentLlmService
from app.core.agent.utils.workflow import WorkflowStep


class GenerateGeneralSummaryAnswerStep(WorkflowStep[GeneralQaSummaryContext]):
    step_id = "generate_general_summary_answer"
    title = "Общий ответ через LLM"

    def __init__(self, llm: AgentLlmService) -> None:
        self._llm = llm

    async def run(self, context: GeneralQaSummaryContext) -> GeneralQaSummaryContext:
        if context.answer:
            return context
        if context.gate_decision is not None and not context.gate_decision.passed:
            context.answer = context.gate_decision.message
            context.answer_generated_payload = {
                "answer_mode": context.gate_decision.answer_mode,
                "answer_length": len(context.answer),
            }
            return context
        if not context.workflow_llm_enabled:
            context.answer = self._build_deterministic_answer(context)
            context.answer_generated_payload = {"answer_mode": "grounded_summary", "answer_length": len(context.answer)}
            return context
        context.prompt_input = self._build_prompt_input(context)
        request_id = context.runtime.request.request_id
        context.answer = await asyncio.to_thread(
            self._llm.generate,
            context.prompt_name,
            context.prompt_input,
            log_context=f"agent:{request_id}",
            trace=context.runtime.trace.module("workflow.v2.general_summary.llm"),
        )
        context.answer_generated_payload = {"answer_mode": "grounded_summary", "answer_length": len(context.answer)}
        return context

    def _build_prompt_input(self, context: GeneralQaSummaryContext) -> str:
        blocks = [f"Запрос пользователя:\n{context.route.user_query}", "Опорные документы:"]
        for index, item in enumerate(context.documents, start=1):
            blocks.append(
                f"{index}. path: {item.path}\n"
                f"title: {item.title}\n"
                f"summary: {item.summary}"
            )
        return "\n\n".join(blocks)

    def _build_deterministic_answer(self, context: GeneralQaSummaryContext) -> str:
        if not context.documents:
            return "В найденной документации нет достаточной опоры для общего summary по запросу."
        return "\n".join(item.summary for item in context.documents[:2] if item.summary)

    def trace_input(self, context: GeneralQaSummaryContext) -> dict[str, object]:
        return {"query": context.route.normalized_query}

    def trace_output(self, context: GeneralQaSummaryContext) -> dict[str, object]:
        return {"answer_length": len(context.answer)}
