from functools import cached_property
from typing import Any, Self

import httpx

from mistralai.workflows import worker_client
from mistralai.workflows.client import get_async_client, translate_model
from mistralai.workflows.exceptions import ErrorCode, WorkflowsException
from mistralai.workflows.models.workflow import WorkflowSpecWithTaskQueue
from mistralai.workflows.protocol.v1.worker import ExecutorIdentityResult, WorkerInfo
from mistralai.workflows.protocol.v1.workflow import (
    DeploymentLocation,
    WorkerHeartbeatResponse,
    WorkflowRegistrationRef,
    WorkflowSpecsRegisterResponse,
)
from mistralai.workflows.worker_client.sdk import PrivateWorkerClient


class _NoopSyncClient:
    """No-op sync HTTP client stub passed to the Speakeasy SDK constructor.

    When ``client=None``, the SDK allocates a real ``httpx.Client`` (opening a
    connection pool) even though ``WorkerClient`` only calls async methods.
    Passing this stub avoids that allocation entirely.
    """

    def send(self, request: httpx.Request, **kwargs: Any) -> httpx.Response:
        raise RuntimeError("Synchronous SDK methods are not supported by WorkerClient")

    def build_request(self, method: str, url: Any, **kwargs: Any) -> httpx.Request:
        return httpx.Request(method, url)

    def close(self) -> None:
        pass


class WorkerClient:
    def __init__(
        self,
        base_url: str,
        timeout: float = 60.0,
        api_key: str | None = None,
        headers: httpx.Headers | dict[str, str] | None = None,
    ):
        """
        Args:
            base_url: The base URL of the Workflows API.
            api_version: The API version to use.
            timeout: The timeout for API requests in seconds.
            api_key: Optional API key for authentication (sent as Bearer token).
        """
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.api_key = api_key
        self._async_client = get_async_client(timeout=timeout, api_key=api_key, headers=headers)

    async def close(self) -> None:
        """Close the underlying HTTP client."""
        await self._async_client.aclose()

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self.close()

    @cached_property
    def private_worker_client(self) -> PrivateWorkerClient:
        return PrivateWorkerClient(
            server_url=self.base_url,
            client=_NoopSyncClient(),
            async_client=self._async_client,
            timeout_ms=int(self.timeout * 1000),
        )

    async def who_am_i(self) -> WorkerInfo:
        try:
            result = await self.private_worker_client.whoami_async()
            return translate_model(WorkerInfo, result)
        except Exception as exc:
            raise WorkflowsException.from_api_client_error(
                exc, message="Failed to get worker info", code=ErrorCode.GET_WORKERS_WHOAMI_ERROR
            ) from exc

    async def fetch_executor_identity(self, execution_token: str) -> ExecutorIdentityResult:
        try:
            result = await self.private_worker_client.executor_identity_async(execution_token=execution_token)
            return translate_model(ExecutorIdentityResult, result)
        except Exception as exc:
            raise WorkflowsException.from_api_client_error(
                exc, message="Failed to get execution identity", code=ErrorCode.POST_EXECUTION_IDENTITY_ERROR
            ) from exc

    async def register_workflow_specs(
        self,
        definitions: list[WorkflowSpecWithTaskQueue],
        deployment_name: str | None = None,
        worker_name: str | None = None,
        deployment_location: DeploymentLocation | None = None,
    ) -> WorkflowSpecsRegisterResponse:
        """Register workflow specs.

        Args:
            definitions: The workflow specs to register.
            deployment_name: Name of the deployment this worker belongs to.
            worker_name: Human-readable name of this worker process (hostname or pod name).
            deployment_location: Metadata about where this deployment is running.
        """
        try:
            result = await self.private_worker_client.register_workflow_definitions_async(
                definitions=[
                    translate_model(worker_client.models.WorkflowSpecWithTaskQueue, definition)
                    for definition in definitions
                ],
                deployment_name=deployment_name,
                worker_name=worker_name,
                deployment_location=translate_model(worker_client.models.DeploymentLocation, deployment_location)
                if deployment_location
                else None,
            )
            return translate_model(WorkflowSpecsRegisterResponse, result)
        except Exception as exc:
            raise WorkflowsException.from_api_client_error(
                exc, message="Failed to register workflow specs", code=ErrorCode.POST_WORKFLOWS_REGISTER_ERROR
            ) from exc

    async def heartbeat(
        self,
        workflow_registration_refs: list[WorkflowRegistrationRef],
        deployment_name: str | None = None,
        worker_name: str | None = None,
    ) -> WorkerHeartbeatResponse:
        try:
            result = await self.private_worker_client.heartbeat_async(
                workflow_registration_refs=[
                    translate_model(worker_client.models.WorkflowRegistrationRefInput, ref)
                    for ref in workflow_registration_refs
                ],
                deployment_name=deployment_name,
                worker_name=worker_name,
            )
            return translate_model(WorkerHeartbeatResponse, result)
        except Exception as exc:
            raise WorkflowsException.from_api_client_error(
                exc, message="Failed to heartbeat", code=ErrorCode.POST_WORKERS_HEARTBEAT_ERROR
            ) from exc
