from collections.abc import AsyncIterator
from uuid import uuid4

from ..utils import generate_good_names
from .shell import JupyterShell
from .transport import InputRequest, KernelMessage


class JupyterApp:
    """Coordinate named persistent shells backed by Jupyter kernels."""

    def __init__(self) -> None:
        self.shells: dict[str, JupyterShell] = {}
        self.last_shell: str | None = None
        self._calls: dict[str, JupyterShell] = {}

    def _new_shell(self) -> tuple[str, JupyterShell]:
        name = generate_good_names()
        while name in self.shells:
            name = generate_good_names()
        shell = JupyterShell(name)
        self.shells[name] = shell
        self.last_shell = name
        return name, shell

    def _select_shell(self, shell_name: str) -> tuple[str, JupyterShell]:
        if shell_name == "new":
            return self._new_shell()
        if shell_name == "last":
            if self.last_shell is None:
                return self._new_shell()
            return self.last_shell, self.shells[self.last_shell]
        if shell_name not in self.shells:
            self.shells[shell_name] = JupyterShell(shell_name)
        self.last_shell = shell_name
        return shell_name, self.shells[shell_name]

    def get_shell(self, shell_name: str) -> JupyterShell:
        """Return a named Jupyter shell."""
        if shell_name not in self.shells:
            raise KeyError(f"Shell {shell_name} does not exist")
        return self.shells[shell_name]

    async def run_code_stream(
        self,
        code: str,
        shell_name: str = "last",
    ) -> AsyncIterator[KernelMessage | InputRequest]:
        """Execute code and yield its subprocess messages."""
        _, shell = self._select_shell(shell_name)
        call_id = uuid4().hex
        self._calls[call_id] = shell
        try:
            async for message in shell.run_cell_stream(code, call_id=call_id):
                yield message
        finally:
            self._calls.pop(call_id, None)

    async def reply_to_input(self, call_id: str, value: str) -> None:
        """Reply to an input request emitted by a running call."""
        try:
            shell = self._calls[call_id]
        except KeyError as error:
            raise KeyError(f"Call {call_id} does not exist") from error
        await shell.reply_to_input(call_id, value)

    async def shutdown(self) -> None:
        """Stop all kernel processes owned by this app."""
        awaitables = [shell.shutdown() for shell in self.shells.values()]
        for awaitable in awaitables:
            await awaitable