from __future__ import annotations

import asyncio
import base64
import json
import logging
from collections.abc import Awaitable, Callable
from datetime import timedelta
from inspect import Parameter, signature
from time import monotonic
from typing import TYPE_CHECKING
from uuid import uuid4

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

from cloud.auth import (
    AuthProvider,
    HostAuthorizationError,
    digest_token,
)
from cloud.internal_api.models import (
    AssignmentModel,
    ClaimRequest,
    ClaimResponse,
    DeviceEnrollmentRequest,
    DeviceEnrollmentResponse,
    HeartbeatRequest,
    HeartbeatResponse,
    HostGovernancePolicyModel,
    HostEnrollmentRequest,
    HostEnrollmentResponse,
    HostTaskCancellationResponse,
    HostTaskSubmissionRequest,
    HostTaskSubmissionResponse,
    LeaseRenewalRequest,
    LeaseRenewalResponse,
    PlannerDecisionError,
    PlannerDecisionRequest,
    PlannerDecisionResponse,
    StaleLeaseConflict,
    TerminalResultRequest,
    TerminalResultResponse,
)
from cloud.llm_providers import LlmProviderResolutionError, LlmProviderService
from cloud.planner_config import build_cloud_planner_client
from cloud.provider_secrets import ProviderSecretConfigurationError
from cloud.repository import (
    AssignmentProgressSnapshot,
    DeviceEnrollmentConflictError,
    HostEnrollmentConflictError,
)
from cloud.governance import TokenBudgetExceededError
from core.models import Device, utc_now
from runtime.tool_calling_client import ToolCallingClient, ToolCallUnavailable
from runtime.tool_specs import ToolSpec

if TYPE_CHECKING:
    from cloud.pool import DevicePool
    from cloud.scheduler import TaskScheduler

logger = logging.getLogger(__name__)


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,
    planner_client_factory: Callable[[], ToolCallingClient] | None = None,
    planner_provider_service: LlmProviderService | None = None,
    scheduler: TaskScheduler | None = None,
    planner_token_reservation_ceiling: int = 4096,
    planner_token_reservation_ttl_seconds: float = 300.0,
) -> 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")
    if planner_token_reservation_ceiling <= 0:
        raise ValueError("planner_token_reservation_ceiling must be greater than zero")
    if planner_token_reservation_ttl_seconds <= 0:
        raise ValueError(
            "planner_token_reservation_ttl_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.post(
        "/enrollments",
        response_model=HostEnrollmentResponse,
        status_code=status.HTTP_201_CREATED,
    )
    def enroll_host(
        payload: HostEnrollmentRequest,
    ) -> HostEnrollmentResponse:
        try:
            enrollment = pool.store.enroll_host(
                host_id=f"host-{uuid4().hex}",
                agent_instance_id=payload.agent_instance_id,
                credential_digest=digest_token(payload.host_token),
                enrollment_token_digest=None,
                display_name=payload.display_name,
                enrolled_at=utc_now(),
            )
        except HostEnrollmentConflictError as exc:
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=str(exc),
            ) from exc
        return HostEnrollmentResponse(host_id=enrollment.host_id)

    @router.post(
        "/hosts/{host_id}/devices/enroll",
        response_model=DeviceEnrollmentResponse,
        status_code=status.HTTP_201_CREATED,
    )
    def enroll_device(
        host_id: str,
        payload: DeviceEnrollmentRequest,
        request: Request,
    ) -> DeviceEnrollmentResponse:
        authorize_host(request, host_id)
        try:
            enrollment = pool.store.enroll_device(
                device_id=f"device-{uuid4().hex}",
                host_id=host_id,
                local_device_id=payload.local_device_id,
                driver_type=payload.driver_type,
                name=payload.name,
                capability_tags=list(payload.capability_tags),
                enrolled_at=utc_now(),
            )
        except DeviceEnrollmentConflictError as exc:
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=str(exc),
            ) from exc
        return DeviceEnrollmentResponse(device_id=enrollment.device_id)

    @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,
            planner_transport=payload.planner_transport,
            mcp_busy_device_ids=payload.mcp_busy_device_ids,
        )
        policy = pool.store.get_host_governance_policy(host_id)
        policy_revision = policy.revision if policy is not None else 0
        return HeartbeatResponse(
            host_id=host_id,
            accepted_devices=len(devices),
            received_at=utc_now(),
            policy_revision=policy_revision,
            policy=(
                HostGovernancePolicyModel(
                    revision=policy.revision,
                    self_submission_enabled=policy.self_submission_enabled,
                    max_active_tasks=policy.max_active_tasks,
                    daily_token_budget=policy.daily_token_budget,
                )
                if policy is not None and payload.policy_revision != policy.revision
                else None
            ),
        )

    @router.post(
        "/hosts/{host_id}/tasks",
        response_model=HostTaskSubmissionResponse,
        status_code=status.HTTP_201_CREATED,
    )
    def submit_host_task(
        host_id: str,
        payload: HostTaskSubmissionRequest,
        request: Request,
    ) -> HostTaskSubmissionResponse:
        authorize_host(request, host_id)
        if payload.host_id != host_id:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail="task host_id must match the request path",
            )
        if scheduler is None:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="task submission is unavailable",
            )
        policy = pool.store.get_host_governance_policy(host_id)
        if policy is not None and not policy.self_submission_enabled:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Host self-submission is disabled",
            )
        if payload.device_id is not None and not any(
            device.host_id == host_id and device.device_id == payload.device_id
            for device in pool.list_devices()
        ):
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail="target device is not owned by this Host",
            )
        from cloud.scheduler import TaskConstraints

        task_id = scheduler.submit(
            goal=payload.goal,
            constraints=TaskConstraints(
                target_host_id=host_id,
                target_device_id=payload.device_id,
            ),
        )
        return HostTaskSubmissionResponse(task_id=task_id)

    @router.post(
        "/hosts/{host_id}/tasks/{task_id}/cancel",
        response_model=HostTaskCancellationResponse,
        responses={
            status.HTTP_202_ACCEPTED: {"model": HostTaskCancellationResponse},
        },
    )
    def cancel_host_task(
        host_id: str,
        task_id: str,
        request: Request,
        response: Response,
    ) -> HostTaskCancellationResponse:
        authorize_host(request, host_id)
        if scheduler is None:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="task cancellation is unavailable",
            )
        task = scheduler.store.get_task(task_id)
        if task is None or task.constraints.target_host_id != host_id:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"task {task_id!r} not found",
            )
        result = scheduler.store.request_task_cancellation(
            task_id, requested_at=utc_now()
        )
        if result == "not_found":
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"task {task_id!r} not found",
            )
        if result == "already_terminal":
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=f"task {task_id!r} has already reached a terminal state",
            )
        task = scheduler.store.get_task(task_id)
        assert task is not None
        if result == "requested" and task.status != "cancelled":
            response.status_code = status.HTTP_202_ACCEPTED
        return HostTaskCancellationResponse(task_id=task_id, status=task.status)

    @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)
        progress_snapshot: AssignmentProgressSnapshot | None = None
        if payload.progress is not None:
            progress_snapshot = AssignmentProgressSnapshot(
                step_index=payload.progress.step_index,
                step_status=payload.progress.step_status,
                summary=payload.progress.summary[:500],
                updated_at=now,
            )
        renewal = 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,
            progress=progress_snapshot,
        )
        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,
            cancel_requested=renewal.cancel_requested,
        )

    @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)

    @router.post(
        "/hosts/{host_id}/planner/decide",
        response_model=PlannerDecisionResponse,
        response_model_exclude_none=True,
        responses={
            status.HTTP_502_BAD_GATEWAY: {"model": PlannerDecisionError},
            status.HTTP_429_TOO_MANY_REQUESTS: {"model": PlannerDecisionError},
        },
    )
    def decide_planner_call(
        host_id: str,
        payload: PlannerDecisionRequest,
        request: Request,
    ):
        authorize_host(request, host_id)
        if payload.host_id != host_id:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail="planner-decision host_id must match the request path",
            )
        _validate_planner_context(pool, host_id=host_id, payload=payload)

        screenshot: bytes | None = None
        if payload.screenshot_base64 is not None:
            try:
                screenshot = base64.b64decode(payload.screenshot_base64)
            except (ValueError, TypeError):
                raise HTTPException(
                    status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                    detail="screenshot_base64 is not valid base64",
                ) from None

        tools = [
            ToolSpec(
                name=tool.name,
                description=tool.description,
                parameters=tool.parameters,
            )
            for tool in payload.tools
        ]

        try:
            if planner_client_factory is not None:
                client = planner_client_factory()
                resolved_provider = "test"
                resolved_model = "test"
                planner_timeout = payload.timeout_seconds
            elif planner_provider_service is not None:
                resolved = planner_provider_service.resolve_active_profile()
                client = build_cloud_planner_client(resolved)
                resolved_provider = resolved.profile.provider_type
                resolved_model = resolved.profile.model
                planner_timeout = resolved.profile.timeout_seconds
            else:
                raise LlmProviderResolutionError(
                    "no database Provider resolver configured"
                )
        except (LlmProviderResolutionError, ProviderSecretConfigurationError) as exc:
            logger.info(
                "planner-decision request failed",
                extra={"host_id": host_id, "error_class": type(exc).__name__},
            )
            return JSONResponse(
                status_code=status.HTTP_502_BAD_GATEWAY,
                content=PlannerDecisionError(detail=str(exc)).model_dump(),
            )

        now = utc_now()
        reservation = None
        try:
            reservation = pool.store.reserve_host_token_budget(
                reservation_id=uuid4().hex,
                host_id=host_id,
                usage_day=now.date().isoformat(),
                reserved_tokens=planner_token_reservation_ceiling,
                task_id=payload.task_id,
                attempt=payload.attempt,
                created_at=now,
                expires_at=now
                + timedelta(seconds=planner_token_reservation_ttl_seconds),
            )
        except TokenBudgetExceededError as exc:
            return JSONResponse(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                content=PlannerDecisionError(detail=str(exc)).model_dump(),
            )

        started_at = monotonic()
        try:
            decision_kwargs = {
                "system_prompt": payload.system_prompt,
                "user_prompt": payload.user_prompt,
                "screenshot": screenshot,
                "tools": tools,
                "timeout": planner_timeout,
            }
            parameters = signature(client.decide).parameters.values()
            if any(
                parameter.name == "history"
                or parameter.kind == Parameter.VAR_KEYWORD
                for parameter in parameters
            ):
                decision_kwargs["history"] = payload.history
            decision = client.decide(**decision_kwargs)
        except ToolCallUnavailable as exc:
            logger.info(
                "planner-decision request failed",
                extra={
                    "host_id": host_id,
                    "latency_seconds": monotonic() - started_at,
                    "error_class": type(exc).__name__,
                },
            )
            return JSONResponse(
                status_code=status.HTTP_502_BAD_GATEWAY,
                content=PlannerDecisionError(detail=str(exc)).model_dump(),
            )
        usage = decision.usage
        if (
            reservation is not None
            and usage is not None
            and usage.total_tokens is not None
        ):
            pool.store.settle_host_token_reservation(
                reservation_id=reservation.id,
                event_id=uuid4().hex,
                provider=resolved_provider,
                model=resolved_model,
                input_tokens=usage.input_tokens,
                output_tokens=usage.output_tokens,
                total_tokens=usage.total_tokens,
                occurred_at=utc_now(),
            )
        if payload.task_id and payload.attempt:
            pool.store.record_planner_decision(
                host_id=host_id,
                task_id=payload.task_id,
                attempt=payload.attempt,
                system_prompt=payload.system_prompt,
                user_prompt=payload.user_prompt,
                tool_name=decision.tool_name,
                arguments_json=json.dumps(decision.arguments),
                now=utc_now(),
                rationale=getattr(decision, "text_output", None),
                thinking=getattr(decision, "thinking", None),
                purpose=getattr(decision, "purpose", None),
                expected_outcome=getattr(decision, "expected_outcome", None),
            )
        logger.info(
            "planner-decision request resolved",
            extra={
                "host_id": host_id,
                "tool_name": decision.tool_name,
                "latency_seconds": monotonic() - started_at,
            },
        )
        return PlannerDecisionResponse(
            tool_name=decision.tool_name,
            arguments=dict(decision.arguments),
            rationale=getattr(decision, "text_output", None),
            thinking=getattr(decision, "thinking", None),
            purpose=getattr(decision, "purpose", None),
            expected_outcome=getattr(decision, "expected_outcome", None),
            input_tokens=(decision.usage.input_tokens if decision.usage else None),
            output_tokens=(decision.usage.output_tokens if decision.usage else None),
            total_tokens=(decision.usage.total_tokens if decision.usage else None),
        )

    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 _validate_planner_context(
    pool, *, host_id: str, payload: PlannerDecisionRequest
) -> None:
    context_values = (payload.task_id, payload.attempt, payload.lease_id)
    if not any(value is not None for value in context_values):
        return
    if any(value is None for value in context_values):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="planner context requires task_id, attempt, and lease_id together",
        )
    attempts = pool.store.list_task_attempts(payload.task_id or "")
    if not any(
        attempt.attempt == payload.attempt
        and attempt.host_id == host_id
        and attempt.lease_id == payload.lease_id
        for attempt in attempts
    ):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="planner context does not match a Host assignment",
        )


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",
        )

    if pool.store.is_enrollment_managed_host(host_id):
        enrollments = {
            enrollment.device_id: enrollment
            for enrollment in pool.store.list_device_enrollments(host_id)
            if enrollment.revoked_at is None
        }
        invalid = [
            device.device_id
            for device in payload.devices
            if device.device_id not in enrollments
            or enrollments[device.device_id].driver_type != device.driver_type
        ]
        if invalid:
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=f"device enrollment conflict: {sorted(set(invalid))}",
            )
        return False

    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