from __future__ import annotations

from typing import Any

from opentelemetry import trace

from terminaluse.lib.core.clients.temporal.temporal_client import TemporalClient
from terminaluse.lib.core.clients.temporal.types import WorkflowState
from terminaluse.lib.core.temporal.types.workflow import SignalName
from terminaluse.lib.environment_variables import EnvironmentVariables
from terminaluse.lib.types.acp import CreateTaskParams, SendEventParams
from terminaluse.lib.utils.logging import make_logger
from terminaluse.types.agent import Agent
from terminaluse.types.event import Event
from terminaluse.types.task import Task

logger = make_logger(__name__)
tracer = trace.get_tracer(__name__)


class TemporalTaskService:
    """
    Submits Agent agent_tasks to the async runtime for execution.
    """

    def __init__(
        self,
        temporal_client: TemporalClient,
        env_vars: EnvironmentVariables,
    ):
        self._temporal_client = temporal_client
        self._env_vars = env_vars

    async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | None) -> str:
        """
        Submit a task to the async runtime for execution.

        returns the workflow ID of the temporal workflow
        """
        with tracer.start_as_current_span(
            "temporal.task_service.submit_task",
            attributes={
                "task.id": task.id,
                "agent.id": agent.id,
                "workflow.name": self._env_vars.TERMINALUSE_WORKFLOW_NAME,
                "workflow.task_queue": self._env_vars.TERMINALUSE_WORKFLOW_TASK_QUEUE,
            },
        ) as span:
            workflow_id = await self._temporal_client.start_workflow(
                workflow=self._env_vars.TERMINALUSE_WORKFLOW_NAME,
                arg=CreateTaskParams(
                    agent=agent,
                    task=task,
                    params=params,
                ),
                id=task.id,
                task_queue=self._env_vars.TERMINALUSE_WORKFLOW_TASK_QUEUE,
            )
            span.set_attribute("workflow.id", workflow_id)
            return workflow_id

    async def get_state(self, task_id: str) -> WorkflowState:
        """
        Get the task state from the async runtime.
        """
        return await self._temporal_client.get_workflow_status(
            workflow_id=task_id,
        )

    async def send_event(self, agent: Agent, task: Task, event: Event, request: dict | None = None) -> None:
        span_attributes: dict[str, str] = {
            "task.id": task.id,
            "agent.id": agent.id,
            "signal.name": SignalName.RECEIVE_EVENT.value,
        }
        if event.id:
            span_attributes["event.id"] = event.id

        with tracer.start_as_current_span(
            "temporal.task_service.send_event",
            attributes=span_attributes,
        ):
            await self._temporal_client.send_signal(
                workflow_id=task.id,
                signal=SignalName.RECEIVE_EVENT.value,
                payload=SendEventParams(
                    agent=agent,
                    task=task,
                    event=event,
                    request=request,
                ).model_dump(),
            )

    async def cancel(self, task_id: str) -> None:
        return await self._temporal_client.cancel_workflow(
            workflow_id=task_id,
        )

    async def terminate(self, task_id: str) -> None:
        return await self._temporal_client.terminate_workflow(
            workflow_id=task_id,
        )
