from __future__ import annotations

import asyncio
import uuid
from collections.abc import Callable
from datetime import UTC, datetime, timedelta

from cloud.internal_api.models import (
    ClaimResponse,
    HeartbeatResponse,
    HostTaskCancellationResponse,
    HostTaskSubmissionResponse,
    LeaseRenewalResponse,
    TerminalResultResponse,
)
from host_agent.progress import TaskProgressSnapshot


class LocalHostAgentClient:
    """In-process task broker used when the Host Agent runs without Cloud."""

    def __init__(self, *, host_id: str, device_ids: Callable[[], list[str]]) -> None:
        self.host_id = host_id
        self._device_ids = device_ids
        self._queue: asyncio.Queue[tuple[str, str, str | None]] = asyncio.Queue()
        self._cancelled: set[str] = set()

    async def submit_self_task(self, *, goal: str, device_id: str | None = None) -> HostTaskSubmissionResponse:
        task_id = f"local-{uuid.uuid4().hex}"
        await self._queue.put((task_id, goal, device_id))
        return HostTaskSubmissionResponse(task_id=task_id)

    async def claim(self):
        task_id, goal, requested_device = await self._queue.get()
        devices = self._device_ids()
        device_id = requested_device or (devices[0] if devices else "")
        if not device_id:
            return None
        from cloud.internal_api.models import AssignmentModel

        return AssignmentModel(
            task_id=task_id,
            attempt=1,
            lease_id=f"local-lease-{uuid.uuid4().hex}",
            lease_expires_at=datetime.now(UTC) + timedelta(days=3650),
            host_id=self.host_id,
            device_id=device_id,
            goal=goal,
        )

    async def heartbeat(self, *args, **kwargs) -> HeartbeatResponse:
        return HeartbeatResponse(
            host_id=self.host_id,
            accepted_devices=len(self._device_ids()),
            received_at=datetime.now(UTC),
        )

    async def renew(self, assignment, *, progress: TaskProgressSnapshot | None = None) -> LeaseRenewalResponse:
        return LeaseRenewalResponse(
            status="renewed",
            lease_expires_at=datetime.now(UTC) + timedelta(days=3650),
        )

    async def report_result(self, assignment, *, status: str, failure_reason: str | None = None, result: dict | None = None) -> TerminalResultResponse:
        return TerminalResultResponse(status="recorded")

    async def cancel_task(self, task_id: str) -> HostTaskCancellationResponse:
        self._cancelled.add(task_id)
        return HostTaskCancellationResponse(task_id=task_id, status="cancel_requested")

    async def aclose(self) -> None:
        return None