import asyncio
import os
import time
from fasta2a.task_manager import TaskManager
from fasta2a.schema import SendMessageResponse

try:
    from fasta2a.schema import SendMessageResult
except ImportError:
    def SendMessageResult(**kwargs):
        return kwargs


class BlockingTaskManager(TaskManager):
    def __init__(self, broker, storage, **kwargs):
        super().__init__(broker, storage, **kwargs)
        self.storage = storage

    async def send_message(self, request):
        response = await super().send_message(request)
        config = request["params"].get("configuration") or {}
        if not config.get("blocking", False):
            return response
        task = response["result"]
        task_id = task["id"]
        timeout = min(
            max(float(os.getenv("A2A_BLOCKING_TIMEOUT", "60")), 1),
            90,
        )
        deadline = time.monotonic() + timeout
        terminal_states = {
            "completed",
            "failed",
            "canceled",
            "rejected",
            "input-required",
            "auth-required",
        }
        while True:
            task = await self.storage.load_task(task_id)
            if task and task["status"]["state"] in terminal_states:
                return SendMessageResponse(
                    jsonrpc="2.0",
                    id=request["id"],
                    result=task,
                )
            if time.monotonic() >= deadline:
                return SendMessageResponse(
                    jsonrpc="2.0",
                    id=request["id"],
                    result=task or response["result"],
                )
            await asyncio.sleep(0.5)