import uuid
from typing import Any, Literal, Protocol

import httpx
from pydantic import BaseModel, Field, ValidationError

from .computer_capabilities import issue_computer_capability
from .config import get_settings
from .errors import AppError


class ComputerWorkerSessionState(BaseModel):
    worker_id: str = Field(min_length=1, max_length=160)
    status: Literal["ready", "ai_controlled", "human_controlled", "paused", "stopped"]
    current_url: str | None = Field(default=None, max_length=4096)
    active_tab_id: str | None = Field(default=None, max_length=160)
    tabs: list[dict[str, Any]] = Field(default_factory=list, max_length=50)
    page_title: str | None = Field(default=None, max_length=500)
    observed_at: str | None = Field(default=None, max_length=80)


class ComputerWorkerActionResult(BaseModel):
    session: ComputerWorkerSessionState
    tool: str = Field(min_length=1, max_length=80)
    result: dict[str, Any]
    verification: dict[str, Any]


class ComputerClient(Protocol):
    async def start_session(
        self, *, workspace_id: uuid.UUID, session_id: uuid.UUID, profile_id: uuid.UUID,
        profile_storage_key: str, fencing_token: int, request_id: str,
    ) -> ComputerWorkerSessionState: ...

    async def command(
        self, *, session_id: uuid.UUID, command: Literal["pause", "resume", "stop"],
        fencing_token: int, request_id: str,
    ) -> ComputerWorkerSessionState: ...

    async def snapshot(
        self, *, session_id: uuid.UUID, fencing_token: int, request_id: str,
    ) -> ComputerWorkerSessionState: ...

    async def action(
        self, *, session_id: uuid.UUID, fencing_token: int, tool: str,
        arguments: dict[str, Any], request_id: str,
    ) -> ComputerWorkerActionResult: ...

    async def human_pointer(
        self, *, session_id: uuid.UUID, fencing_token: int, payload: dict[str, Any],
        request_id: str,
    ) -> ComputerWorkerSessionState: ...

    async def human_keyboard(
        self, *, session_id: uuid.UUID, fencing_token: int, payload: dict[str, Any],
        request_id: str,
    ) -> ComputerWorkerSessionState: ...

    async def frame(
        self, *, session_id: uuid.UUID, fencing_token: int, request_id: str,
    ) -> dict[str, Any]: ...

    async def artifact(
        self, *, workspace_id: uuid.UUID, session_id: uuid.UUID, artifact_id: uuid.UUID,
        request_id: str,
    ) -> bytes: ...

    async def delete_artifact(
        self, *, workspace_id: uuid.UUID, session_id: uuid.UUID, artifact_id: uuid.UUID,
        request_id: str,
    ) -> None: ...


class HttpComputerClient:
    @staticmethod
    def _headers(
        *, scope: str, request_id: str, session_id: uuid.UUID | None = None,
        workspace_id: uuid.UUID | None = None, fencing_token: int | None = None,
    ) -> dict[str, str]:
        settings = get_settings()
        token = settings.read_computer_service_token()
        if not token:
            raise AppError(
                "COMPUTER_SERVICE_AUTH_UNCONFIGURED",
                "Computer service authentication is unconfigured.",
                503,
            )
        capability = issue_computer_capability(
            scope=scope, request_id=request_id, session_id=session_id,
            workspace_id=workspace_id, fencing_token=fencing_token,
        )
        return {
            "x-service-token": token,
            "x-request-id": request_id,
            "x-computer-capability": capability,
        }

    async def _request(
        self, method: str, path: str, *, request_id: str, scope: str,
        session_id: uuid.UUID | None = None, workspace_id: uuid.UUID | None = None,
        fencing_token: int | None = None, payload: dict | None = None,
    ) -> ComputerWorkerSessionState:
        settings = get_settings()
        try:
            async with httpx.AsyncClient(timeout=30) as client:
                response = await client.request(
                    method,
                    f"{settings.computer_service_url.rstrip('/')}{path}",
                    headers=self._headers(
                        scope=scope, request_id=request_id, session_id=session_id,
                        workspace_id=workspace_id, fencing_token=fencing_token,
                    ),
                    json=payload,
                )
        except httpx.HTTPError as error:
            raise AppError(
                "COMPUTER_SERVICE_UNAVAILABLE", "The computer service is unavailable.", 503
            ) from error
        if not response.is_success:
            try:
                body = response.json()
                service_error = body["error"]
                code = str(service_error["code"])
                message = str(service_error["message"])
            except (ValueError, KeyError, TypeError):
                code, message = (
                    "COMPUTER_SERVICE_FAILED",
                    "The computer service could not complete the request.",
                )
            safe_status = response.status_code if response.status_code in {
                401, 404, 409, 422, 503
            } else 502
            raise AppError(code[:80], message[:500], safe_status)
        try:
            return ComputerWorkerSessionState.model_validate(response.json())
        except (ValueError, ValidationError) as error:
            raise AppError(
                "COMPUTER_SERVICE_RESPONSE_INVALID",
                "The computer service returned an invalid response.",
                502,
            ) from error

    async def start_session(
        self, *, workspace_id: uuid.UUID, session_id: uuid.UUID, profile_id: uuid.UUID,
        profile_storage_key: str, fencing_token: int, request_id: str,
    ) -> ComputerWorkerSessionState:
        return await self._request(
            "POST", "/v1/sessions", request_id=request_id, scope="session:start",
            session_id=session_id, workspace_id=workspace_id, fencing_token=fencing_token,
            payload={
                "workspace_id": str(workspace_id),
                "session_id": str(session_id),
                "profile_id": str(profile_id),
                "profile_storage_key": profile_storage_key,
                "fencing_token": fencing_token,
            },
        )

    async def command(
        self, *, session_id: uuid.UUID, command: Literal["pause", "resume", "stop"],
        fencing_token: int, request_id: str,
    ) -> ComputerWorkerSessionState:
        return await self._request(
            "POST", f"/v1/sessions/{session_id}/{command}", request_id=request_id,
            scope=f"session:{command}", session_id=session_id, fencing_token=fencing_token,
            payload={"fencing_token": fencing_token},
        )

    async def snapshot(
        self, *, session_id: uuid.UUID, fencing_token: int, request_id: str,
    ) -> ComputerWorkerSessionState:
        return await self._request(
            "POST", f"/v1/sessions/{session_id}/snapshot", request_id=request_id,
            scope="session:snapshot", session_id=session_id, fencing_token=fencing_token,
            payload={"fencing_token": fencing_token},
        )

    async def action(
        self, *, session_id: uuid.UUID, fencing_token: int, tool: str,
        arguments: dict[str, Any], request_id: str,
    ) -> ComputerWorkerActionResult:
        settings = get_settings()
        try:
            async with httpx.AsyncClient(timeout=40) as client:
                response = await client.post(
                    f"{settings.computer_service_url.rstrip('/')}/v1/sessions/"
                    f"{session_id}/actions",
                    headers=self._headers(
                        scope=f"browser:{tool}", request_id=request_id,
                        session_id=session_id, fencing_token=fencing_token,
                    ),
                    json={
                        "fencing_token": fencing_token,
                        "tool": tool,
                        "arguments": arguments,
                    },
                )
        except httpx.HTTPError as error:
            raise AppError(
                "COMPUTER_SERVICE_UNAVAILABLE", "The computer service is unavailable.", 503
            ) from error
        if not response.is_success:
            self._raise_service_error(response)
        try:
            return ComputerWorkerActionResult.model_validate(response.json())
        except (ValueError, ValidationError) as error:
            raise AppError(
                "COMPUTER_SERVICE_RESPONSE_INVALID",
                "The computer service returned an invalid response.", 502,
            ) from error

    @staticmethod
    def _raise_service_error(response: httpx.Response) -> None:
        try:
            service_error = response.json()["error"]
            code = str(service_error["code"])
            message = str(service_error["message"])
        except (ValueError, KeyError, TypeError):
            code, message = (
                "COMPUTER_SERVICE_FAILED",
                "The computer service could not complete the request.",
            )
        safe_status = response.status_code if response.status_code in {
            401, 403, 404, 409, 422, 503, 504
        } else 502
        raise AppError(code[:80], message[:500], safe_status)

    async def _human_request(
        self, *, session_id: uuid.UUID, path: str, fencing_token: int,
        payload: dict[str, Any], request_id: str,
    ) -> ComputerWorkerSessionState:
        settings = get_settings()
        try:
            async with httpx.AsyncClient(timeout=35) as client:
                response = await client.post(
                    f"{settings.computer_service_url.rstrip('/')}/v1/sessions/"
                    f"{session_id}/{path}",
                    headers=self._headers(
                        scope=f"human:{path.split('/')[-1]}", request_id=request_id,
                        session_id=session_id, fencing_token=fencing_token,
                    ),
                    json={"fencing_token": fencing_token, **payload},
                )
        except httpx.HTTPError as error:
            raise AppError(
                "COMPUTER_SERVICE_UNAVAILABLE", "The computer service is unavailable.", 503
            ) from error
        if not response.is_success:
            self._raise_service_error(response)
        try:
            return ComputerWorkerSessionState.model_validate(response.json())
        except (ValueError, ValidationError) as error:
            raise AppError(
                "COMPUTER_SERVICE_RESPONSE_INVALID",
                "The computer service returned an invalid response.", 502,
            ) from error

    async def human_pointer(
        self, *, session_id: uuid.UUID, fencing_token: int, payload: dict[str, Any],
        request_id: str,
    ) -> ComputerWorkerSessionState:
        return await self._human_request(
            session_id=session_id, path="human/pointer", fencing_token=fencing_token,
            payload=payload, request_id=request_id,
        )

    async def human_keyboard(
        self, *, session_id: uuid.UUID, fencing_token: int, payload: dict[str, Any],
        request_id: str,
    ) -> ComputerWorkerSessionState:
        return await self._human_request(
            session_id=session_id, path="human/keyboard", fencing_token=fencing_token,
            payload=payload, request_id=request_id,
        )

    async def frame(
        self, *, session_id: uuid.UUID, fencing_token: int, request_id: str,
    ) -> dict[str, Any]:
        settings = get_settings()
        try:
            async with httpx.AsyncClient(timeout=20) as client:
                response = await client.post(
                    f"{settings.computer_service_url.rstrip('/')}/v1/sessions/"
                    f"{session_id}/frame",
                    headers=self._headers(
                        scope="session:frame", request_id=request_id,
                        session_id=session_id, fencing_token=fencing_token,
                    ),
                    json={"fencing_token": fencing_token},
                )
        except httpx.HTTPError as error:
            raise AppError(
                "COMPUTER_SERVICE_UNAVAILABLE", "The computer service is unavailable.", 503
            ) from error
        if not response.is_success:
            self._raise_service_error(response)
        try:
            payload = response.json()
        except ValueError as error:
            raise AppError(
                "COMPUTER_SERVICE_RESPONSE_INVALID",
                "The computer service returned an invalid response.", 502,
            ) from error
        if not isinstance(payload, dict) or payload.get("mime_type") != "image/jpeg":
            raise AppError(
                "COMPUTER_SERVICE_RESPONSE_INVALID",
                "The computer service returned an invalid response.", 502,
            )
        return payload

    async def artifact(
        self, *, workspace_id: uuid.UUID, session_id: uuid.UUID, artifact_id: uuid.UUID,
        request_id: str,
    ) -> bytes:
        settings = get_settings()
        try:
            async with httpx.AsyncClient(timeout=20) as client:
                response = await client.get(
                    f"{settings.computer_service_url.rstrip('/')}/v1/artifacts/"
                    f"{workspace_id}/{session_id}/{artifact_id}",
                    headers=self._headers(
                        scope="artifact:read", request_id=request_id,
                        session_id=session_id, workspace_id=workspace_id,
                    ),
                )
        except httpx.HTTPError as error:
            raise AppError(
                "COMPUTER_SERVICE_UNAVAILABLE", "The computer service is unavailable.", 503
            ) from error
        if not response.is_success:
            self._raise_service_error(response)
        if response.headers.get("content-type", "").split(";", 1)[0] != "image/jpeg":
            raise AppError(
                "COMPUTER_ARTIFACT_INVALID", "The browser artifact is invalid.", 502
            )
        if len(response.content) > 20_000_000:
            raise AppError(
                "COMPUTER_ARTIFACT_INVALID", "The browser artifact is invalid.", 502
            )
        return response.content

    async def delete_artifact(
        self, *, workspace_id: uuid.UUID, session_id: uuid.UUID, artifact_id: uuid.UUID,
        request_id: str,
    ) -> None:
        settings = get_settings()
        try:
            async with httpx.AsyncClient(timeout=20) as client:
                response = await client.delete(
                    f"{settings.computer_service_url.rstrip('/')}/v1/artifacts/"
                    f"{workspace_id}/{session_id}/{artifact_id}",
                    headers=self._headers(
                        scope="artifact:delete", request_id=request_id,
                        session_id=session_id, workspace_id=workspace_id,
                    ),
                )
        except httpx.HTTPError as error:
            raise AppError(
                "COMPUTER_SERVICE_UNAVAILABLE", "The computer service is unavailable.", 503
            ) from error
        if response.status_code not in {204, 404}:
            self._raise_service_error(response)


def get_computer_client() -> ComputerClient:
    return HttpComputerClient()
