# This file was auto-generated by Fern from our API Definition.

import base64
import json
import typing
from json.decoder import JSONDecodeError
from pathlib import Path

import websockets
import websockets.sync.connection as websockets_sync_connection

from ....core.api_error import ApiError
from ....core.events import EventEmitterMixin, EventType
from ....core.pydantic_utilities import parse_obj_as
from .types.config import Config
from .types.stream_models_endpoint_payload import StreamModelsEndpointPayload
from .types.subscribe_event import SubscribeEvent

try:
    from websockets.legacy.client import WebSocketClientProtocol  # type: ignore
except ImportError:
    from websockets import WebSocketClientProtocol  # type: ignore

StreamSocketClientResponse = typing.Union[SubscribeEvent]


class AsyncStreamSocketClient(EventEmitterMixin):
    def __init__(self, *, websocket: WebSocketClientProtocol):
        super().__init__()
        self._websocket = websocket

    async def __aiter__(self):
        async for message in self._websocket:
            yield parse_obj_as(StreamSocketClientResponse, json.loads(message))  # type: ignore

    async def start_listening(self):
        """
        Start listening for messages on the websocket connection.

        Emits events in the following order:
        - EventType.OPEN when connection is established
        - EventType.MESSAGE for each message received
        - EventType.ERROR if an error occurs
        - EventType.CLOSE when connection is closed
        """
        await self._emit_async(EventType.OPEN, None)
        try:
            async for raw_message in self._websocket:
                json_data = json.loads(raw_message)
                parsed = parse_obj_as(StreamSocketClientResponse, json_data)  # type: ignore
                await self._emit_async(EventType.MESSAGE, parsed)
        except (websockets.WebSocketException, JSONDecodeError) as exc:
            await self._emit_async(EventType.ERROR, exc)
        finally:
            await self._emit_async(EventType.CLOSE, None)

    async def send_publish(self, message: StreamModelsEndpointPayload) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a StreamModelsEndpointPayload.
        """
        await self._send_model(message)

    async def recv(self) -> StreamSocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = await self._websocket.recv()
        json_data = json.loads(data)
        return parse_obj_as(StreamSocketClientResponse, json_data)  # type: ignore

    async def _send(self, data: typing.Any) -> None:
        """
        Send a message to the websocket connection.
        """
        if isinstance(data, dict):
            data = json.dumps(data)
        await self._websocket.send(data)

    async def _send_model(self, data: typing.Any) -> None:
        """
        Send a Pydantic model to the websocket connection.
        """
        await self._send(data.dict())

    async def send_facemesh(
        self,
        landmarks: typing.List[typing.List[typing.List[float]]],
        config: typing.Optional[Config] = None,
        payload_id: typing.Optional[str] = None,
    ) -> StreamSocketClientResponse:
        landmarks_str = json.dumps(landmarks)
        payload = {
            "data": landmarks_str,
            "models": config.dict() if config else None,
            "raw_text": False,
            "payload_id": payload_id,
        }
        payload = {k: v for k, v in payload.items() if v is not None}
        await self._websocket.send(json.dumps(payload))
        return await self.recv()

    async def send_text(
        self,
        text: str,
        config: typing.Optional[Config] = None,
        payload_id: typing.Optional[str] = None,
    ) -> StreamSocketClientResponse:
        payload = {
            "data": text,
            "models": config.dict() if config else None,
            "raw_text": True,
            "payload_id": payload_id,
        }
        payload = {k: v for k, v in payload.items() if v is not None}
        await self._websocket.send(json.dumps(payload))
        return await self.recv()

    async def send_file(
        self,
        file_: typing.Union[str, Path],
        config: typing.Optional[Config] = None,
        payload_id: typing.Optional[str] = None,
    ) -> StreamSocketClientResponse:
        try:
            with open(file_, "rb") as f:
                bytes_data = base64.b64encode(f.read()).decode()
        except:
            if isinstance(file_, Path):
                raise ApiError(body=f"Failed to open file: {file_}")
            # If you cannot open the file, assume you were passed a b64 string, not a file path
            bytes_data = str(file_)

        payload = {
            "data": bytes_data,
            "models": config.dict() if config else None,
            "raw_text": False,
            "payload_id": payload_id,
        }
        payload = {k: v for k, v in payload.items() if v is not None}
        await self._websocket.send(json.dumps(payload))
        return await self.recv()

    async def get_job_details(self) -> StreamSocketClientResponse:
        payload = {"job_details": True}
        await self._websocket.send(json.dumps(payload))
        return await self.recv()

    async def reset(self) -> StreamSocketClientResponse:
        payload = {"reset_stream": True}
        await self._websocket.send(json.dumps(payload))
        return await self.recv()


class StreamSocketClient(EventEmitterMixin):
    def __init__(self, *, websocket: websockets_sync_connection.Connection):
        super().__init__()
        self._websocket = websocket

    def __iter__(self):
        for message in self._websocket:
            yield parse_obj_as(StreamSocketClientResponse, json.loads(message))  # type: ignore

    def start_listening(self):
        """
        Start listening for messages on the websocket connection.

        Emits events in the following order:
        - EventType.OPEN when connection is established
        - EventType.MESSAGE for each message received
        - EventType.ERROR if an error occurs
        - EventType.CLOSE when connection is closed
        """
        self._emit(EventType.OPEN, None)
        try:
            for raw_message in self._websocket:
                json_data = json.loads(raw_message)
                parsed = parse_obj_as(StreamSocketClientResponse, json_data)  # type: ignore
                self._emit(EventType.MESSAGE, parsed)
        except (websockets.WebSocketException, JSONDecodeError) as exc:
            self._emit(EventType.ERROR, exc)
        finally:
            self._emit(EventType.CLOSE, None)

    def send_publish(self, message: StreamModelsEndpointPayload) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a StreamModelsEndpointPayload.
        """
        self._send_model(message)

    def recv(self) -> StreamSocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = self._websocket.recv()
        json_data = json.loads(data)
        return parse_obj_as(StreamSocketClientResponse, json_data)  # type: ignore

    def _send(self, data: typing.Any) -> None:
        """
        Send a message to the websocket connection.
        """
        if isinstance(data, dict):
            data = json.dumps(data)
        self._websocket.send(data)

    def _send_model(self, data: typing.Any) -> None:
        """
        Send a Pydantic model to the websocket connection.
        """
        self._send(data.dict())

    def send_facemesh(
        self,
        landmarks: typing.List[typing.List[typing.List[float]]],
        config: typing.Optional[Config] = None,
        payload_id: typing.Optional[str] = None,
    ) -> StreamSocketClientResponse:
        landmarks_str = json.dumps(landmarks)
        payload = {
            "data": landmarks_str,
            "models": config.dict() if config else None,
            "raw_text": False,
            "payload_id": payload_id,
        }
        payload = {k: v for k, v in payload.items() if v is not None}
        self._websocket.send(json.dumps(payload))
        return self.recv()

    def send_text(
        self,
        text: str,
        config: typing.Optional[Config] = None,
        payload_id: typing.Optional[str] = None,
    ) -> StreamSocketClientResponse:
        payload = {
            "data": text,
            "models": config.dict() if config else None,
            "raw_text": True,
            "payload_id": payload_id,
        }
        payload = {k: v for k, v in payload.items() if v is not None}
        self._websocket.send(json.dumps(payload))
        return self.recv()

    def send_file(
        self,
        file_: typing.Union[str, Path],
        config: typing.Optional[Config] = None,
        payload_id: typing.Optional[str] = None,
    ) -> StreamSocketClientResponse:
        try:
            with open(file_, "rb") as f:
                bytes_data = base64.b64encode(f.read()).decode()
        except:
            if isinstance(file_, Path):
                raise ApiError(body=f"Failed to open file: {file_}")
            # If you cannot open the file, assume you were passed a b64 string, not a file path
            bytes_data = str(file_)

        payload = {
            "data": bytes_data,
            "models": config.dict() if config else None,
            "raw_text": False,
            "payload_id": payload_id,
        }
        payload = {k: v for k, v in payload.items() if v is not None}
        self._websocket.send(json.dumps(payload))
        return self.recv()

    def get_job_details(self) -> StreamSocketClientResponse:
        payload = {"job_details": True}
        self._websocket.send(json.dumps(payload))
        return self.recv()

    def reset(self) -> StreamSocketClientResponse:
        payload = {"reset_stream": True}
        self._websocket.send(json.dumps(payload))
        return self.recv()
