from __future__ import annotations

from app.modules.agent.orchestration.adapters.llm_chat_adapter import LlmChatAdapter
from app.modules.agent.orchestration.context.execution_context import ExecutionContext
from app.modules.agent.orchestration.processes.v1.prompt_payload_builder import V1PromptPayloadBuilder


class SimpleLlmWorkflow:
    workflow_id = "simple_llm"

    def __init__(
        self,
        llm: LlmChatAdapter,
        prompt_payload_builder: V1PromptPayloadBuilder | None = None,
    ) -> None:
        self._llm = llm
        self._payload_builder = prompt_payload_builder or V1PromptPayloadBuilder()

    async def run(self, context: ExecutionContext) -> dict:
        workflow_trace = context.trace.module("task_workflow")
        workflow_trace.log("started", {"workflow_id": self.workflow_id})
        prompt_payload = self._payload_builder.build(context)
        answer = await self._llm.generate(
            prompt_payload,
            context.request.request_id,
            trace=context.trace.module("llm"),
        )
        result = {
            "workflow_id": self.workflow_id,
            "prompt_name": self._llm.prompt_name,
            "prompt_payload": prompt_payload,
            "answer": answer,
        }
        workflow_trace.log(
            "completed",
            {
                "workflow_id": self.workflow_id,
                "prompt_name": self._llm.prompt_name,
                "answer_length": len(answer),
            },
        )
        return result
