from datetime import UTC, datetime, timedelta

from sqlalchemy import select, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.ext.asyncio import AsyncSession

from app.model.platform import AgentRun


class AgentRunRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get(self, run_id: str) -> AgentRun | None:
        result = await self.session.execute(select(AgentRun).where(AgentRun.run_id == run_id))
        return result.scalar_one_or_none()

    async def claim(self, run_id: str, worker_id: str, lease_seconds: int) -> bool:
        now = datetime.now(UTC).replace(tzinfo=None)
        result = await self.session.execute(
            update(AgentRun)
            .where(AgentRun.run_id == run_id, AgentRun.status.in_(("queued", "running")))
            .where((AgentRun.locked_until.is_(None)) | (AgentRun.locked_until < now))
            .values(status="running", worker_id=worker_id,
                    locked_until=now + timedelta(seconds=lease_seconds),
                    attempt_count=AgentRun.attempt_count + 1)
        )
        return isinstance(result, CursorResult) and result.rowcount == 1

    async def request_cancel(self, run_id: str, when: datetime) -> bool:
        result = await self.session.execute(
            update(AgentRun)
            .where(AgentRun.run_id == run_id)
            .where(AgentRun.status.in_(("queued", "running")))
            .values(status="cancel_requested", cancel_requested_at=when)
        )
        return isinstance(result, CursorResult) and result.rowcount == 1

    async def renew(self, run_id: str, worker_id: str, lease_seconds: int) -> bool:
        now = datetime.now(UTC).replace(tzinfo=None)
        result = await self.session.execute(
            update(AgentRun)
            .where(AgentRun.run_id == run_id, AgentRun.worker_id == worker_id,
                   AgentRun.status == "running", AgentRun.locked_until > now)
            .values(locked_until=now + timedelta(seconds=lease_seconds))
        )
        return isinstance(result, CursorResult) and result.rowcount == 1

    async def is_cancel_requested(self, run_id: str) -> bool:
        run = await self.get(run_id)
        return run is not None and run.status == "cancel_requested"

    async def record_failure(self, run_id: str, error_code: str, retry_limit: int) -> bool:
        run = await self.get(run_id)
        if run is None:
            return False
        if run.attempt_count < retry_limit:
            run.status = "queued"
            run.error_code = error_code
            run.locked_until = None
            await self.session.flush()
            return True
        run.status = "failed"
        run.error_code = error_code
        run.locked_until = None
        await self.session.flush()
        return False