# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

"""
Shared logic for client-side handling of subscriptions to the
Pyre daemon. Context on this:
- A Pyre daemon opens a socket accepting client connections.
- The socket accepts one-off requests using a request/response model
- But the socket also allows long-running connections where we send a
  subscription request, and the connection stays open. The server will
  then push notifications (type errors and/or state change alerts) to
  the client.
"""

import dataclasses
import json
from typing import Optional, Union

from . import incremental

TypeErrors = incremental.TypeErrors


def _parse_type_error_subscription(response: object) -> TypeErrors:
    return incremental.parse_type_error_response_json(["TypeErrors", response])


@dataclasses.dataclass(frozen=True)
class StatusUpdate:
    kind: str
    message: Optional[str] = None


def _parse_status_update_subscription(response: object) -> StatusUpdate:
    if not isinstance(response, list) or len(response) == 0:
        raise incremental.InvalidServerResponse(
            f"Status update subscription must be a nonempty list. Got {response}"
        )
    kind = response[0]
    if not isinstance(kind, str):
        raise incremental.InvalidServerResponse(
            f"Response kind of a status update must be a string. Got {response}"
        )
    message = None
    if len(response) > 1 and isinstance(response[1], dict) and "message" in response[1]:
        message = response[1]["message"]
    return StatusUpdate(kind=kind, message=message)


@dataclasses.dataclass(frozen=True)
class Error:
    message: str


def _parse_error_subscription(response: object) -> Error:
    if not isinstance(response, str):
        raise incremental.InvalidServerResponse(
            f"Response kind of an error must be a string. Got {response}"
        )
    return Error(message=response)


class IncrementalTelemetry:
    pass


Body = Union[TypeErrors, StatusUpdate, Error, IncrementalTelemetry]


@dataclasses.dataclass(frozen=True)
class Response:
    body: Body

    @staticmethod
    def parse(response: str) -> "Response":
        try:
            response_json = json.loads(response)
            # The response JSON is expected to have the following forms:
            # `{"name": "foo", "body": ["TypeErrors", [error_json, ...]]}`
            # `{"name": "foo", "body": ["StatusUpdate", ["message_kind", ...]]}`
            # `{"name": "foo", "body": ["Error", "error message"]}`
            # The legacy subscription protocol includes a subscription name - we have
            # since moved away from requiring a name associated with a subscription, as
            # there aren't practical ways to set up multiple subscriptions on the same connection,
            # and a connection serves as a unique identifier for a subscription.

            # For backwards compatibility, we allow the "name" field to be in the response, but do not
            # parse it on the client.
            if isinstance(response_json, dict):
                body = response_json.get("body", None)
                if body is not None and isinstance(body, list) and len(body) > 1:
                    tag = body[0]
                    if tag == "TypeErrors":
                        return Response(body=_parse_type_error_subscription(body[1]))
                    elif tag == "StatusUpdate":
                        return Response(body=_parse_status_update_subscription(body[1]))
                    elif tag == "Error":
                        return Response(body=_parse_error_subscription(body[1]))
                    elif tag == "IncrementalTelemetry":
                        return Response(body=IncrementalTelemetry())

            raise incremental.InvalidServerResponse(
                f"Unexpected JSON subscription from server: {response_json}"
            )
        except json.JSONDecodeError as decode_error:
            message = f"Cannot parse subscription as JSON: {decode_error}"
            raise incremental.InvalidServerResponse(message) from decode_error
