# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import urllib.parse
from typing import TYPE_CHECKING
from unittest import mock

import tornado.testing
import tornado.web
import tornado.websocket

from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
from streamlit.runtime import Runtime
from streamlit.runtime.app_session import AppSession
from streamlit.web.server import Server

if TYPE_CHECKING:
    from tornado.websocket import WebSocketClientConnection


class ServerTestCase(tornado.testing.AsyncHTTPTestCase):
    """Base class for async streamlit.server testing.

    Subclasses should patch 'streamlit.server.server.AppSession',
    to prevent AppSessions from being created, and scripts from
    being run. (Script running involves creating new threads, which
    interfere with other tests if not properly terminated.)

    See the "ServerTest" class for example usage.
    """

    _next_session_id = 0

    def tearDown(self) -> None:
        super().tearDown()
        # Server._create_app() will create the Runtime singleton.
        # We null it out in tearDown() so that it doesn't interfere with
        # future tests.
        Runtime._instance = None

    def get_app(self) -> tornado.web.Application:
        self.server = Server(
            "/not/a/script.py",
            is_hello=False,
        )
        app = self.server._create_app()
        # Store app settings for tests.
        self.app_settings = app.settings
        return app

    def get_ws_url(self, path):
        """Return a ws:// URL with the given path for our test server."""
        # get_url() gives us a result with the 'http' scheme;
        # we swap it out for 'ws'.
        url = self.get_url(path)
        parts = list(urllib.parse.urlparse(url))
        parts[0] = "ws"
        return urllib.parse.urlunparse(tuple(parts))

    async def ws_connect(self, existing_session_id=None) -> WebSocketClientConnection:
        """Open a websocket connection to the server.

        Returns
        -------
        WebSocketClientConnection
            The connected websocket client.
        """
        # See the comment in WebsocketConnection.tsx about how we repurpose the
        # Sec-WebSocket-Protocol header for more information on how this works.
        if existing_session_id is None:
            subprotocols = ["streamlit", "PLACEHOLDER_AUTH_TOKEN"]
        else:
            subprotocols = ["streamlit", "PLACEHOLDER_AUTH_TOKEN", existing_session_id]

        return await tornado.websocket.websocket_connect(
            self.get_ws_url("/_stcore/stream"),
            subprotocols=subprotocols,
        )

    async def read_forward_msg(
        self, ws_client: WebSocketClientConnection
    ) -> ForwardMsg:
        """Parse the next message from a Websocket client into a ForwardMsg."""
        data = await ws_client.read_message()
        message = ForwardMsg()
        message.ParseFromString(data)
        return message

    @staticmethod
    def _create_mock_app_session(*args, **kwargs):
        """Create a mock AppSession. Each mocked instance will have
        its own unique ID."""
        mock_id = mock.PropertyMock(
            return_value=f"mock_id:{ServerTestCase._next_session_id}"
        )
        ServerTestCase._next_session_id += 1

        mock_session = mock.MagicMock(AppSession, autospec=True, *args, **kwargs)  # noqa: B026
        type(mock_session).id = mock_id
        return mock_session

    def _patch_app_session(self):
        """Mock the Server's AppSession import. We don't want
        actual sessions to be instantiated, or scripts to be run.
        """

        return mock.patch(
            "streamlit.runtime.websocket_session_manager.AppSession",
            # new_callable must return a function, not an object, or else
            # there will only be a single AppSession mock. Hence the lambda.
            new_callable=lambda: self._create_mock_app_session,
        )
