from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from time import monotonic
from typing import TYPE_CHECKING

from fastapi import APIRouter, HTTPException, Request, status

from cloud.auth import (
    AuthProvider,
    HostAuthorizationError,
)
from cloud.internal_api.models import (
    AssignmentModel,
    ClaimRequest,
    ClaimResponse,
    HeartbeatRequest,
    HeartbeatResponse,
)
from core.models import Device, utc_now

if TYPE_CHECKING:
    from cloud.pool import DevicePool


def create_internal_router(
    *,
    pool: DevicePool,
    auth_provider: AuthProvider,
    version_prefix: str = "/internal/v1",
    claim_poll_interval_seconds: float = 0.1,
    sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
) -> APIRouter:
    if claim_poll_interval_seconds <= 0:
        raise ValueError("claim_poll_interval_seconds must be greater than zero")
    router = APIRouter(prefix=version_prefix, tags=["host-agent"])

    def authorize_host(request: Request, host_id: str) -> None:
        principal = auth_provider.authenticate(request)
        if principal is None:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="unauthorized",
                headers={"WWW-Authenticate": "Bearer"},
            )
        try:
            principal.require_host(host_id)
        except HostAuthorizationError as exc:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=str(exc),
            ) from exc

    @router.put(
        "/hosts/{host_id}/heartbeat",
        response_model=HeartbeatResponse,
    )
    def heartbeat(
        host_id: str,
        payload: HeartbeatRequest,
        request: Request,
    ) -> HeartbeatResponse:
        authorize_host(request, host_id)
        _validate_snapshot(pool, host_id=host_id, payload=payload)
        devices = [
            Device(
                id=device.device_id,
                driver_type=device.driver_type,
                status=device.status,
                capability_tags=list(device.capability_tags),
            )
            for device in payload.devices
        ]
        pool.sync_host_devices(host_id, devices, address=payload.address)
        return HeartbeatResponse(
            host_id=host_id,
            accepted_devices=len(devices),
            received_at=utc_now(),
        )

    @router.post(
        "/hosts/{host_id}/assignments/claim",
        response_model=ClaimResponse,
    )
    async def claim_assignment(
        host_id: str,
        payload: ClaimRequest,
        request: Request,
    ) -> ClaimResponse:
        authorize_host(request, host_id)
        if payload.host_id != host_id:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail="claim host_id must match the request path",
            )

        deadline = monotonic() + payload.timeout_seconds
        while True:
            assignment = pool.store.claim_assignment(host_id=host_id, now=utc_now())
            if assignment is not None:
                return ClaimResponse(
                    assignment=AssignmentModel(
                        task_id=assignment.task_id,
                        attempt=assignment.attempt,
                        lease_id=assignment.lease_id,
                        lease_expires_at=assignment.lease_expires_at,
                        host_id=assignment.host_id,
                        device_id=assignment.device_id,
                        goal=assignment.goal,
                        workflow_definition_id=assignment.workflow_definition_id,
                    )
                )

            remaining = deadline - monotonic()
            if remaining <= 0:
                return ClaimResponse(timed_out=True)
            await sleep(min(claim_poll_interval_seconds, remaining))

    return router


def _validate_snapshot(
    pool: DevicePool,
    *,
    host_id: str,
    payload: HeartbeatRequest,
) -> None:
    if payload.host_id != host_id:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="heartbeat host_id must match the request path",
        )
    device_ids = [device.device_id for device in payload.devices]
    if len(device_ids) != len(set(device_ids)):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="heartbeat snapshot contains duplicate device ids",
        )

    now = utc_now()
    hosts = {host.host_id: host for host in pool.store.list_hosts()}
    conflicts: list[str] = []
    requested_ids = set(device_ids)
    for device in pool.store.list_devices():
        if device.device_id not in requested_ids or device.host_id == host_id:
            continue
        owner = hosts.get(device.host_id)
        if owner is None:
            continue
        age_seconds = (now - owner.last_seen_at).total_seconds()
        if age_seconds <= pool.config.stale_after_seconds:
            conflicts.append(device.device_id)
    if conflicts:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail=f"device ownership conflict: {sorted(set(conflicts))}",
        )