from __future__ import annotations

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

from fastapi import APIRouter, HTTPException, Request, status
from fastapi.responses import JSONResponse

from cloud.auth import (
    AuthProvider,
    HostAuthorizationError,
)
from cloud.internal_api.models import (
    AssignmentModel,
    ClaimRequest,
    ClaimResponse,
    HeartbeatRequest,
    HeartbeatResponse,
    LeaseRenewalRequest,
    LeaseRenewalResponse,
    StaleLeaseConflict,
    TerminalResultRequest,
    TerminalResultResponse,
)
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,
    lease_duration_seconds: float = 60.0,
    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")
    if lease_duration_seconds <= 0:
        raise ValueError("lease_duration_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)
        allow_device_takeover = _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,
            allow_device_takeover=allow_device_takeover,
        )
        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))

    @router.post(
        "/hosts/{host_id}/assignments/{task_id}/renew",
        response_model=LeaseRenewalResponse,
        responses={status.HTTP_409_CONFLICT: {"model": StaleLeaseConflict}},
    )
    def renew_assignment(
        host_id: str,
        task_id: str,
        payload: LeaseRenewalRequest,
        request: Request,
    ):
        authorize_host(request, host_id)
        _validate_assignment_identity(
            host_id=host_id,
            task_id=task_id,
            payload_host_id=payload.host_id,
            payload_task_id=payload.task_id,
        )
        now = utc_now()
        lease_expires_at = now + timedelta(seconds=lease_duration_seconds)
        renewal_status = pool.store.renew_lease(
            task_id=task_id,
            attempt=payload.attempt,
            lease_id=payload.lease_id,
            host_id=host_id,
            lease_expires_at=lease_expires_at,
            now=now,
        )
        if renewal_status == "not_found":
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="assignment not found",
            )
        if renewal_status != "renewed":
            return _stale_lease_conflict("assignment lease is stale or expired")
        return LeaseRenewalResponse(
            status="renewed",
            lease_expires_at=lease_expires_at,
        )

    @router.post(
        "/hosts/{host_id}/assignments/{task_id}/result",
        response_model=TerminalResultResponse,
        responses={status.HTTP_409_CONFLICT: {"model": StaleLeaseConflict}},
    )
    def report_result(
        host_id: str,
        task_id: str,
        payload: TerminalResultRequest,
        request: Request,
    ):
        authorize_host(request, host_id)
        _validate_assignment_identity(
            host_id=host_id,
            task_id=task_id,
            payload_host_id=payload.host_id,
            payload_task_id=payload.task_id,
        )
        result_status = pool.store.record_task_result(
            task_id=task_id,
            attempt=payload.attempt,
            lease_id=payload.lease_id,
            host_id=host_id,
            status=payload.status,
            failure_reason=payload.failure_reason,
            terminal_result=payload.result,
            completed_at=utc_now(),
        )
        if result_status == "conflict":
            return _stale_lease_conflict("assignment lease is stale or superseded")
        return TerminalResultResponse(status=result_status)

    return router


def _validate_assignment_identity(
    *,
    host_id: str,
    task_id: str,
    payload_host_id: str,
    payload_task_id: str,
) -> None:
    if payload_host_id != host_id or payload_task_id != task_id:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="assignment identity must match the request path",
        )


def _stale_lease_conflict(detail: str) -> JSONResponse:
    return JSONResponse(
        status_code=status.HTTP_409_CONFLICT,
        content=StaleLeaseConflict(detail=detail).model_dump(),
    )


def _validate_snapshot(
    pool: DevicePool,
    *,
    host_id: str,
    payload: HeartbeatRequest,
) -> bool:
    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] = []
    stale_owner_found = False
    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)
        else:
            stale_owner_found = True
    if conflicts:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail=f"device ownership conflict: {sorted(set(conflicts))}",
        )
    return stale_owner_found