from collections.abc import Awaitable, Callable

from sqlalchemy.ext.asyncio import AsyncSession

from app.core.errors import AgentError, RecoverableAgentError
from app.repository.agent_run_repository import AgentRunRepository


class AgentRunWorker:
    def __init__(
        self,
        session: AsyncSession,
        execute_run: Callable[[str], Awaitable[None]],
        worker_id: str,
        lease_seconds: int = 60,
        retry_limit: int = 3,
    ) -> None:
        self.repository = AgentRunRepository(session)
        self.execute_run = execute_run
        self.worker_id = worker_id
        self.lease_seconds = lease_seconds
        self.retry_limit = retry_limit

    async def claim_and_execute(self, run_id: str) -> bool:
        async with self.repository.session.begin():
            claimed = await self.repository.claim(run_id, self.worker_id, self.lease_seconds)
        if not claimed:
            return False
        try:
            await self.execute_run(run_id)
        except RecoverableAgentError as exc:
            async with self.repository.session.begin():
                await self.repository.record_failure(run_id, exc.code, self.retry_limit)
            return False
        except AgentError as exc:
            async with self.repository.session.begin():
                await self.repository.record_failure(run_id, exc.code, 0)
            return False
        return True

    async def renew(self, run_id: str) -> bool:
        return await self.repository.renew(run_id, self.worker_id, self.lease_seconds)

    async def cancellation_requested(self, run_id: str) -> bool:
        return await self.repository.is_cancel_requested(run_id)