import asyncio
import unittest

from ipython_shell.jupyter_app import JupyterApp
from ipython_shell.jupyter_shell import JupyterShell
from ipython_shell.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_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.assertTrue(input_request.call_id)
            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()