# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import asyncio
import json
import os
import selectors
import signal
import struct
import sys
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, TypedDict

from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState

from marimo import _loggers
from marimo._server.api.auth import validate_auth
from marimo._server.api.deps import AppState
from marimo._server.codes import WebSocketCloseReason, WebSocketCodes
from marimo._server.router import APIRouter
from marimo._session.model import SessionMode
from marimo._utils.asyncio_utils import cancel_and_wait
from marimo._utils.platform import is_pyodide, is_windows

if TYPE_CHECKING:
    from collections.abc import Callable, Iterable

LOGGER = _loggers.marimo_logger()

router = APIRouter()


# Configuration constants
MAX_CHUNK_SIZE = 8192
MAX_BUFFER_SIZE = 65536
MAX_COMMAND_BUFFER_SIZE = 1024
KEEP_COMMAND_CHARS = 512
HEALTH_CHECK_INTERVAL = 5.0
READ_TIMEOUT = 0.05
IDLE_SLEEP = 0.01


def _resize_pty(fd: int, rows: int, cols: int) -> None:
    """Resize the PTY to the specified dimensions."""
    try:
        # Use TIOCSWINSZ ioctl to set window size
        import fcntl
        import termios

        # Format: struct winsize { unsigned short ws_row, ws_col, ws_xpixel, ws_ypixel }
        winsize = struct.pack("HHHH", rows, cols, 0, 0)
        fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
        LOGGER.debug(f"PTY resized to {cols}x{rows}")
    except Exception as e:
        LOGGER.warning(f"Failed to resize PTY: {e}")


def _send_sigwinch(pid: int) -> None:
    """Send SIGWINCH signal to the child process to notify of terminal resize."""
    try:
        os.kill(pid, signal.SIGWINCH)
        LOGGER.debug(f"Sent SIGWINCH to process {pid}")
    except (OSError, ProcessLookupError) as e:
        LOGGER.debug(f"Failed to send SIGWINCH to {pid}: {e}")


def _create_shell_environment(
    cwd: str | None = None,
) -> tuple[str, dict[str, str]]:
    """Create a proper shell environment with working directory and env vars."""
    # Set working directory to a reasonable default
    if cwd is None:
        try:
            cwd = os.getcwd()
        except (OSError, PermissionError):
            try:
                cwd = str(Path.home())
            except Exception:
                cwd = "/tmp"

    # Set up environment variables
    env = os.environ.copy()
    env["TERM"] = "xterm-256color"
    env["LANG"] = env.get("LANG", "en_US.UTF-8")
    env["LC_ALL"] = env.get("LC_ALL", "en_US.UTF-8")

    # Determine shell
    default_shell = os.environ.get("SHELL")
    if not default_shell or not os.path.exists(default_shell):
        # Try common shells
        for shell in ["/bin/bash", "/bin/zsh", "/bin/sh"]:
            if os.path.exists(shell):
                default_shell = shell
                break
        else:
            default_shell = "/bin/sh"  # Fallback

    return default_shell, env


def _setup_child_process(shell: str, env: dict[str, str], cwd: str) -> None:
    """Set up the child process environment and start the shell."""
    try:
        os.chdir(cwd)
        os.execve(shell, [shell], env)
    except Exception as e:
        LOGGER.error(f"Failed to start shell: {e}")
        sys.exit(1)


def _create_process_cleanup_handler(
    child_pid: int, fd: int
) -> Callable[[], None]:
    """Create a cleanup handler for the child process and file descriptor."""

    def cleanup() -> None:
        # Close the pty master *first*.
        #
        # child_pid is the session leader of the pty. If it is killed while the
        # server still holds the master fd open -- and a foreground child (e.g.
        # a coding agent or REPL) still holds the slave -- the kernel cannot
        # finish revoking the shell's controlling terminal, so the shell never
        # becomes a reapable zombie. The blocking os.waitpid() below would then
        # hang the asyncio event loop forever (frozen server, Ctrl-C ignored).
        # Closing the master is the hangup that lets the teardown complete, so
        # the child reaps promptly.
        try:
            os.close(fd)
        except OSError:
            pass

        try:
            # Try graceful termination first
            os.kill(child_pid, signal.SIGTERM)
            # Give the process a moment to terminate gracefully
            try:
                os.waitpid(child_pid, os.WNOHANG)
            except (OSError, ChildProcessError):
                pass
            # Force kill if still running
            try:
                os.kill(child_pid, signal.SIGKILL)
                os.waitpid(child_pid, 0)
            except (OSError, ProcessLookupError, ChildProcessError):
                pass
        except Exception as e:
            LOGGER.debug(f"Error during cleanup: {e}")

    return cleanup


def _decode_pty_data(
    buffer: bytes, max_buffer_size: int = MAX_BUFFER_SIZE
) -> tuple[str, bytes]:
    """Decode PTY data handling partial UTF-8 sequences."""
    try:
        # Decode and send data, handling partial UTF-8 sequences
        text = buffer.decode("utf-8", errors="ignore")
        return text, b""
    except UnicodeDecodeError:
        # Keep buffer if we have incomplete UTF-8 sequence
        if len(buffer) > max_buffer_size:
            # Force send to prevent memory issues
            text = buffer.decode("utf-8", errors="replace")
            return text, b""
        return "", buffer


def _should_close_on_command(command_buffer: str, data: str) -> bool:
    """Check if the terminal should close based on the command."""
    if data in ["\r", "\n"]:
        return command_buffer.strip().lower() == "exit"
    return False


def _manage_command_buffer(
    buffer: str, data: str, max_size: int = MAX_COMMAND_BUFFER_SIZE
) -> str:
    """Manage the command buffer size to prevent memory issues."""
    buffer += data
    if len(buffer) > max_size:
        return buffer[-KEEP_COMMAND_CHARS:]  # Keep last chars
    return buffer


def _wait_for_read(master: int, timeout: float) -> bool:
    """Wait for data to be available on the file descriptor.

    Uses selectors instead of select.select() to avoid the FD_SETSIZE
    limitation (fd must be < 1024) that causes failures when many file
    descriptors are open.
    """
    sel = selectors.DefaultSelector()
    try:
        sel.register(master, selectors.EVENT_READ)
        events = sel.select(timeout=timeout)
        return len(events) > 0
    finally:
        sel.close()


async def _read_from_pty(master: int, websocket: WebSocket) -> None:
    """Read data from PTY and send to websocket with proper buffering.

    Uses raw os.read() instead of os.fdopen() to avoid double-close issues
    when both reader and writer share the same fd.
    """
    loop = asyncio.get_running_loop()
    buffer = b""

    try:
        while True:
            try:
                # Check for available data with a timeout
                has_data = await loop.run_in_executor(
                    None,
                    _wait_for_read,
                    master,
                    READ_TIMEOUT,
                )

                if has_data:
                    # Read available data
                    try:
                        chunk = os.read(master, MAX_CHUNK_SIZE)
                        if not chunk:
                            break
                        buffer += chunk
                    except OSError as e:
                        if e.errno == 5:  # Input/output error (process died)
                            break
                        raise

                # Send buffered data if we have any
                if buffer:
                    text, buffer = _decode_pty_data(buffer)
                    if text:
                        await websocket.send_text(text)
                else:
                    # Small delay to prevent busy-waiting when no data
                    await asyncio.sleep(IDLE_SLEEP)

            except (asyncio.CancelledError, WebSocketDisconnect):
                break
    except OSError as e:
        if e.errno == 9:  # Bad file descriptor
            LOGGER.debug("File descriptor closed, stopping read loop")
            return
        raise  # Re-raise other OSErrors


class ResizeMessage(TypedDict):
    type: Literal["resize"]
    cols: int
    rows: int


async def _maybe_handle_resize(
    *, master: int, child_pid: int, message: str
) -> bool:
    """Handle resize messages from websocket."""

    try:
        parsed_message: ResizeMessage = json.loads(message)
        if (
            isinstance(parsed_message, dict)
            and parsed_message.get("type") == "resize"
        ):
            cols = parsed_message.get("cols")
            rows = parsed_message.get("rows")
            if (
                isinstance(cols, int)
                and isinstance(rows, int)
                and cols > 0
                and rows > 0
            ):
                _resize_pty(master, rows, cols)
                _send_sigwinch(child_pid)
                return True
            else:
                LOGGER.warning("Invalid resize message")
                return False
    except (json.JSONDecodeError, TypeError):
        # Not a JSON message, treat as regular terminal input
        pass

    return False


async def _write_to_pty(
    master: int, websocket: WebSocket, child_pid: int
) -> None:
    """Write data from websocket to PTY with command monitoring and resize handling.

    Uses raw os.write() instead of os.fdopen() to avoid double-close issues
    when both reader and writer share the same fd.
    """
    try:
        command_buffer = ""
        while True:
            try:
                data = await websocket.receive_text()
                LOGGER.debug("Received: %s", repr(data))

                # Check if this is a resize message
                if await _maybe_handle_resize(
                    master=master, child_pid=child_pid, message=data
                ):
                    continue

                # Handle special key combinations and commands
                command_buffer = _manage_command_buffer(command_buffer, data)

                # Check for exit command
                if _should_close_on_command(command_buffer, data):
                    LOGGER.debug("Exit command received, closing connection")
                    return

                # Reset buffer on line endings
                if data in ["\r", "\n"]:
                    command_buffer = ""

                # Write data to PTY
                try:
                    encoded_data = data.encode("utf-8")
                    os.write(master, encoded_data)
                except OSError as e:
                    if e.errno == 5:  # Input/output error (process died)
                        LOGGER.debug("Process died, stopping write loop")
                        break
                    raise

            except (asyncio.CancelledError, WebSocketDisconnect):
                break
    except OSError as e:
        if e.errno == 9:  # Bad file descriptor
            LOGGER.debug("File descriptor closed, stopping write loop")
            return
        raise


async def _cancel_tasks(tasks: Iterable[asyncio.Task[Any]]) -> None:
    for task in tasks:
        await cancel_and_wait(task)


def supports_terminal() -> bool:
    """Whether the current environment supports terminals."""
    if is_windows() or is_pyodide():
        return False
    # Try to import pty to verify it's actually available
    try:
        import pty  # noqa: F401

        return True
    except ImportError:
        return False


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket) -> None:
    app_state = AppState(websocket)

    if app_state.enable_auth and not validate_auth(websocket):
        await websocket.close(
            WebSocketCodes.UNAUTHORIZED, WebSocketCloseReason.UNAUTHORIZED
        )
        return

    if app_state.mode != SessionMode.EDIT:
        await websocket.close(
            code=1008, reason="Terminal only available in edit mode"
        )
        return

    if not supports_terminal():
        await websocket.close(
            code=1008, reason="Terminal not supported in this environment"
        )
        return

    try:
        await websocket.accept()
        LOGGER.debug("Terminal websocket accepted")
    except Exception as e:
        LOGGER.error(f"Failed to accept websocket connection: {e}")
        return

    # Initialize PTY - wrap in try-except to handle any unexpected failures
    try:
        import pty

        # TODO(akshayka): Someone should clean this up to make it safe on
        # macOS.
        #
        # NOTE: On macOS, pty.fork() is documented as unsafe when mixed with
        # higher-level system APIs, and forking a multi-threaded process (as
        # the server is) can segfault the child before it reaches execve().
        # See https://docs.python.org/3/library/pty.html
        child_pid, fd = pty.fork()
        if child_pid == 0:
            # Child process - set up the shell environment
            shell, env = _create_shell_environment()
            cwd = env.get("PWD", os.getcwd())
            _setup_child_process(shell, env, cwd)

        # Set up cleanup handler
        cleanup_child = _create_process_cleanup_handler(child_pid, fd)
    except Exception as e:
        LOGGER.error(f"Failed to initialize terminal: {e}")
        try:
            if websocket.application_state != WebSocketState.DISCONNECTED:
                await websocket.close(
                    code=1011, reason="Failed to initialize terminal"
                )
        except Exception:
            pass
        return

    reader_task = asyncio.create_task(_read_from_pty(fd, websocket))
    writer_task = asyncio.create_task(_write_to_pty(fd, websocket, child_pid))

    try:
        _done, pending = await asyncio.wait(
            [reader_task, writer_task], return_when=asyncio.FIRST_COMPLETED
        )

        # Cancel all pending tasks
        await _cancel_tasks(pending)

    except WebSocketDisconnect:
        LOGGER.debug("Client disconnected from terminal")
    except Exception as e:
        LOGGER.exception(f"Terminal websocket error: {e}")
    finally:
        # Ensure all tasks are cleaned up
        await _cancel_tasks([reader_task, writer_task])

        # Close websocket if still connected
        try:
            if websocket.application_state != WebSocketState.DISCONNECTED:
                await websocket.close(
                    code=1000, reason="Terminal session ended"
                )
        except Exception as e:
            LOGGER.debug(f"Error closing websocket: {e}")

        # Clean up process and file descriptor
        cleanup_child()
    LOGGER.debug("Terminal websocket closed")
