import asyncio
import re
import unittest

from ipython_shell.jupyter.app import JupyterApp
from ipython_shell.jupyter.shell import JupyterShell
from ipython_shell.jupyter.transport import InputRequest, KernelMessage


class JupyterShellTests(unittest.IsolatedAsyncioTestCase):
    async def test_shell_keeps_namespace_between_cells(self):
        shell = JupyterShell("persistent")
        try:
            first = [
                message
                async for message in shell.run_cell_stream(
                    "answer = 41; answer",
                    call_id="call-1",
                )
            ]
            second = [
                message
                async for message in shell.run_cell_stream(
                    "answer + 1",
                    call_id="call-2",
                )
            ]
        finally:
            await shell.shutdown()

        self.assertEqual(
            next(
                message.content["data"]["text/plain"]
                for message in first
                if isinstance(message, KernelMessage)
                and message.msg_type == "execute_result"
            ),
            "41",
        )
        self.assertEqual(
            next(
                message.content["data"]["text/plain"]
                for message in second
                if isinstance(message, KernelMessage)
                and message.msg_type == "execute_result"
            ),
            "42",
        )


class JupyterAppTests(unittest.IsolatedAsyncioTestCase):
    async def test_app_exposes_shared_shell_metadata_api(self):
        app = JupyterApp()
        try:
            created = app.create_shell()
            listed = app.list_shells()
            last = app.get_last_shell()
            selected = app.get_shell(created.name)
            info = app.info()
        finally:
            await app.shutdown()

        self.assertEqual(len(listed), 1)
        self.assertEqual(created.shell_id, listed[0].shell_id)
        self.assertEqual(created.shell_id, last.shell_id)
        self.assertEqual(created.shell_id, selected.shell_id)
        self.assertEqual(created.execution_count, 0)
        self.assertEqual(info.shell_count, 1)
        self.assertEqual(info.last_shell, created.name)

    async def test_named_shells_have_isolated_namespaces(self):
        app = JupyterApp()
        try:
            _alpha_setup = [
                message
                async for message in app.run_code_stream(
                    "alpha_only = 'alpha'",
                    shell_name="alpha",
                )
            ]
            beta_view = [
                message
                async for message in app.run_code_stream(
                    "'alpha_only' in globals()",
                    shell_name="beta",
                )
            ]
            alpha_view = [
                message
                async for message in app.run_code_stream(
                    "alpha_only",
                    shell_name="alpha",
                )
            ]
        finally:
            await app.shutdown()

        self.assertEqual(
            next(
                message.content["data"]["text/plain"]
                for message in beta_view
                if isinstance(message, KernelMessage)
                and message.msg_type == "execute_result"
            ),
            "False",
        )
        self.assertEqual(
            next(
                message.content["data"]["text/plain"]
                for message in alpha_view
                if isinstance(message, KernelMessage)
                and message.msg_type == "execute_result"
            ),
            "'alpha'",
        )

    async def test_app_routes_input_reply_to_the_call(self):
        app = JupyterApp()
        stream = app.run_code_stream(
            "answer = input('name? '); answer",
            shell_name="new",
        )

        try:
            while True:
                event = await asyncio.wait_for(anext(stream), timeout=5)
                if isinstance(event, InputRequest):
                    input_request = event
                    break

            self.assertIsInstance(input_request, InputRequest)
            assert isinstance(input_request, InputRequest)
            self.assertRegex(
                input_request.call_id,
                re.compile(r"^[a-z]+-[a-z]+-[a-z]+-[0-9a-f]{4}$"),
            )
            self.assertEqual(input_request.prompt, "name? ")

            await app.reply_to_input(input_request.call_id, "Ada")
            messages = [message async for message in stream]
        finally:
            await app.shutdown()

        self.assertEqual(
            next(
                message.content["data"]["text/plain"]
                for message in messages
                if isinstance(message, KernelMessage)
                and message.msg_type == "execute_result"
            ),
            "'Ada'",
        )


if __name__ == "__main__":
    unittest.main()