"""
Mock kernel backend for testing.

Provides an in-memory kernel simulation that actually executes Python code
without any external dependencies (no Jupyter Server required).

Features:
- Variable persistence across executions per kernel
- stdout capture as stream outputs
- Expression results as execute_result outputs
- Error output for exceptions
- Input request handling via threading coordination

Usage:
    backend = MockKernelBackend()
    await backend.initialize()
    kernel_id = await backend.start_kernel()
    async for output in backend.execute_code(kernel_id, "print('hello')"):
        print(output)
"""

import asyncio
import contextlib
import io
import threading
import traceback
import uuid
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Dict, List, Optional

from .kernel_backend import (
    KernelBackend,
    KernelConnectionInfo,
    KernelInfo,
    KernelInputError,
    KernelNotFoundError,
    KernelStatus,
    KernelTimeoutError,
    register_backend,
)


@dataclass
class MockKernel:
    """A simulated kernel with its own Python namespace."""

    kernel_id: str
    kernel_name: str = "python3"
    python_path: Optional[str] = None
    env_name: Optional[str] = None
    execution_count: int = 0
    namespace: Dict[str, Any] = field(default_factory=lambda: {"__builtins__": __builtins__})
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    # Input handling (threading-based for sync exec compatibility)
    _input_event: threading.Event = field(default_factory=threading.Event)
    _input_value: Optional[str] = field(default=None)
    _awaiting_input: bool = field(default=False)
    _input_prompt: str = field(default="")


class MockKernelBackend:
    """
    In-memory mock kernel backend for testing.

    Executes Python code via exec/eval with per-kernel namespaces.
    Implements the full KernelBackend protocol without external dependencies.
    """

    def __init__(self, **kwargs):
        self._kernels: Dict[str, MockKernel] = {}
        self._next_id = 0
        self._initialized = False
        self._root_dir: Optional[str] = None

    # =========================================================================
    # Lifecycle Hooks
    # =========================================================================

    async def initialize(self, **kwargs: Any) -> None:
        self._root_dir = kwargs.get("root_dir")
        self._initialized = True

    async def close(self) -> None:
        await self.shutdown_all()
        self._initialized = False

    # =========================================================================
    # Backend State
    # =========================================================================

    def is_ready(self) -> bool:
        return self._initialized

    async def set_root_directory(self, path: str) -> str:
        self._root_dir = path
        return path

    # =========================================================================
    # Kernel Lifecycle
    # =========================================================================

    async def start_kernel(
        self,
        kernel_name: str = "python3",
        python_path: Optional[str] = None,
        env_name: Optional[str] = None,
        cwd: Optional[str] = None,
    ) -> str:
        kernel_id = f"mock-kernel-{self._next_id}"
        self._next_id += 1
        self._kernels[kernel_id] = MockKernel(
            kernel_id=kernel_id,
            kernel_name=kernel_name,
            python_path=python_path,
            env_name=env_name,
        )
        return kernel_id

    async def stop_kernel(self, kernel_id: str) -> None:
        if kernel_id not in self._kernels:
            raise KernelNotFoundError(f"Kernel {kernel_id} not found")
        del self._kernels[kernel_id]

    async def restart_kernel(self, kernel_id: str) -> None:
        kernel = self._get_kernel(kernel_id)
        kernel.namespace = {"__builtins__": __builtins__}
        kernel.execution_count = 0

    async def interrupt_kernel(self, kernel_id: str) -> None:
        self._get_kernel(kernel_id)  # Verify exists

    async def shutdown_all(self, timeout: float = 5.0) -> Dict[str, Any]:
        stopped = list(self._kernels.keys())
        self._kernels.clear()
        return {"stopped": stopped, "failed": {}}

    # =========================================================================
    # Code Execution
    # =========================================================================

    async def execute_code(
        self,
        kernel_id: str,
        code: str,
        timeout: float = 300.0,
    ) -> AsyncIterator[dict]:
        kernel = self._get_kernel(kernel_id)

        async with kernel.lock:
            kernel.execution_count += 1

            # Fast path: no input() in code — run synchronously
            if "input(" not in code:
                for output in self._exec_sync(kernel, code):
                    yield output
                return

            # Slow path: input() support via threading
            async for output in self._exec_with_input(kernel, code, timeout):
                yield output

    def _exec_sync(self, kernel: MockKernel, code: str) -> List[dict]:
        """Execute code synchronously, return list of outputs."""
        outputs = []
        stdout = io.StringIO()

        try:
            with contextlib.redirect_stdout(stdout):
                try:
                    result = eval(code, kernel.namespace)
                    text = stdout.getvalue()
                    if text:
                        outputs.append({"output_type": "stream", "name": "stdout", "text": text})
                    if result is not None:
                        kernel.namespace["_"] = result
                        outputs.append({
                            "output_type": "execute_result",
                            "data": {"text/plain": repr(result)},
                            "metadata": {},
                            "execution_count": kernel.execution_count,
                        })
                except SyntaxError:
                    # Not an expression — exec as statements
                    stdout = io.StringIO()
                    with contextlib.redirect_stdout(stdout):
                        exec(code, kernel.namespace)
                    text = stdout.getvalue()
                    if text:
                        outputs.append({"output_type": "stream", "name": "stdout", "text": text})
        except Exception as e:
            tb = traceback.format_exception(type(e), e, e.__traceback__)
            outputs.append({
                "output_type": "error",
                "ename": type(e).__name__,
                "evalue": str(e),
                "traceback": tb,
            })

        return outputs

    async def _exec_with_input(
        self, kernel: MockKernel, code: str, timeout: float
    ) -> AsyncIterator[dict]:
        """Execute code that may call input(), using threading for coordination."""
        outputs: List[dict] = []
        exec_complete = threading.Event()

        def mock_input(prompt=""):
            kernel._input_prompt = prompt
            kernel._awaiting_input = True
            kernel._input_event.clear()
            if not kernel._input_event.wait(timeout=timeout):
                kernel._awaiting_input = False
                raise TimeoutError("Input timed out")
            kernel._awaiting_input = False
            return kernel._input_value or ""

        kernel.namespace["input"] = mock_input

        def run_code():
            stdout = io.StringIO()
            try:
                with contextlib.redirect_stdout(stdout):
                    try:
                        result = eval(code, kernel.namespace)
                        text = stdout.getvalue()
                        if text:
                            outputs.append({"output_type": "stream", "name": "stdout", "text": text})
                        if result is not None:
                            kernel.namespace["_"] = result
                            outputs.append({
                                "output_type": "execute_result",
                                "data": {"text/plain": repr(result)},
                                "metadata": {},
                                "execution_count": kernel.execution_count,
                            })
                    except SyntaxError:
                        stdout = io.StringIO()
                        with contextlib.redirect_stdout(stdout):
                            exec(code, kernel.namespace)
                        text = stdout.getvalue()
                        if text:
                            outputs.append({"output_type": "stream", "name": "stdout", "text": text})
            except Exception as e:
                tb = traceback.format_exception(type(e), e, e.__traceback__)
                outputs.append({
                    "output_type": "error",
                    "ename": type(e).__name__,
                    "evalue": str(e),
                    "traceback": tb,
                })
            finally:
                exec_complete.set()

        thread = threading.Thread(target=run_code, daemon=True)
        thread.start()

        # Poll for input requests and completion
        while not exec_complete.is_set():
            if kernel._awaiting_input:
                yield {
                    "output_type": "input_request",
                    "prompt": kernel._input_prompt,
                    "password": False,
                }
                # Wait for send_input_reply to unblock the thread
                while kernel._awaiting_input and not exec_complete.is_set():
                    await asyncio.sleep(0.01)
            await asyncio.sleep(0.01)

        # Yield all collected outputs
        for output in outputs:
            yield output

    async def send_input_reply(self, kernel_id: str, value: str) -> None:
        kernel = self._get_kernel(kernel_id)
        if not kernel._awaiting_input:
            raise KernelInputError(f"Kernel {kernel_id} is not waiting for input")
        kernel._input_value = value
        kernel._input_event.set()

    # =========================================================================
    # Status and Discovery
    # =========================================================================

    async def get_kernel_status(self, kernel_id: str) -> Dict[str, Any]:
        if kernel_id in self._kernels:
            return {"status": "running", "kernel_id": kernel_id}
        return {"status": "not_found", "kernel_id": kernel_id}

    def get_kernel_info(self, kernel_id: str) -> Optional[KernelInfo]:
        kernel = self._kernels.get(kernel_id)
        if kernel is None:
            return None
        return KernelInfo(
            kernel_id=kernel.kernel_id,
            status=KernelStatus.IDLE,
            python_path=kernel.python_path,
            env_name=kernel.env_name,
            execution_count=kernel.execution_count,
        )

    def get_kernel_connection_info(self, kernel_id: str) -> Optional[KernelConnectionInfo]:
        kernel = self._kernels.get(kernel_id)
        if kernel is None:
            return None
        return KernelConnectionInfo(
            session_id=f"mock-session-{kernel.kernel_id}",
            python_path=kernel.python_path,
            env_name=kernel.env_name,
        )

    def list_kernels(self) -> List[Dict[str, Any]]:
        return [
            {
                "kernel_id": k.kernel_id,
                "status": "running",
                "python_path": k.python_path,
                "env_name": k.env_name,
            }
            for k in self._kernels.values()
        ]

    async def list_available_kernels(self) -> List[Dict[str, Any]]:
        return [{"name": "python3", "display_name": "Python 3 (Mock)", "language": "python"}]

    def get_execution_count(self, kernel_id: str) -> int:
        return self._get_kernel(kernel_id).execution_count

    async def validate_kernel(self, kernel_id: str) -> bool:
        return kernel_id in self._kernels

    # =========================================================================
    # Internal
    # =========================================================================

    def _get_kernel(self, kernel_id: str) -> MockKernel:
        if kernel_id not in self._kernels:
            raise KernelNotFoundError(f"Kernel {kernel_id} not found")
        return self._kernels[kernel_id]


# Register with the backend factory
register_backend("mock", MockKernelBackend)
