autocoder/server/routers/terminal.py

"""
Terminal Router
===============

REST and WebSocket endpoints for interactive terminal I/O with PTY support.
Provides real-time bidirectional communication with terminal sessions.
Supports multiple terminals per project with create, list, rename, delete operations.
"""

import asyncio
import base64
import json
import logging
import re

from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect
from pydantic import BaseModel

from ..services.terminal_manager import (
    create_terminal,
    delete_terminal,
    get_terminal_info,
    get_terminal_session,
    list_terminals,
    rename_terminal,
    stop_terminal_session,
)
from ..utils.project_helpers import get_project_path as _get_project_path
from ..utils.validation import is_valid_project_name as validate_project_name

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/terminal", tags=["terminal"])


class TerminalCloseCode:
    """WebSocket close codes for terminal endpoint."""

    INVALID_PROJECT_NAME = 4000
    PROJECT_NOT_FOUND = 4004
    FAILED_TO_START = 4500


def validate_terminal_id(terminal_id: str) -> bool:
    """
    Validate terminal ID format.

    Args:
        terminal_id: The terminal ID to validate

    Returns:
        True if valid, False otherwise
    """
    return bool(re.match(r"^[a-zA-Z0-9]{1,16}$", terminal_id))


# Pydantic models for request/response bodies
class CreateTerminalRequest(BaseModel):
    """Request body for creating a terminal."""

    name: str | None = None


class RenameTerminalRequest(BaseModel):
    """Request body for renaming a terminal."""

    name: str


class TerminalInfoResponse(BaseModel):
    """Response model for terminal info."""

    id: str
    name: str
    created_at: str


# REST Endpoints


@router.get("/{project_name}")
async def list_project_terminals(project_name: str) -> list[TerminalInfoResponse]:
    """
    List all terminals for a project.

    Args:
        project_name: Name of the project

    Returns:
        List of terminal info objects
    """
    if not validate_project_name(project_name):
        raise HTTPException(status_code=400, detail="Invalid project name")

    project_dir = _get_project_path(project_name)
    if not project_dir:
        raise HTTPException(status_code=404, detail="Project not found")

    terminals = list_terminals(project_name)

    # If no terminals exist, create a default one
    if not terminals:
        info = create_terminal(project_name)
        terminals = [info]

    return [
        TerminalInfoResponse(id=t.id, name=t.name, created_at=t.created_at) for t in terminals
    ]


@router.post("/{project_name}")
async def create_project_terminal(
    project_name: str, request: CreateTerminalRequest
) -> TerminalInfoResponse:
    """
    Create a new terminal for a project.

    Args:
        project_name: Name of the project
        request: Request body with optional terminal name

    Returns:
        The created terminal info
    """
    if not validate_project_name(project_name):
        raise HTTPException(status_code=400, detail="Invalid project name")

    project_dir = _get_project_path(project_name)
    if not project_dir:
        raise HTTPException(status_code=404, detail="Project not found")

    info = create_terminal(project_name, request.name)
    return TerminalInfoResponse(id=info.id, name=info.name, created_at=info.created_at)


@router.patch("/{project_name}/{terminal_id}")
async def rename_project_terminal(
    project_name: str, terminal_id: str, request: RenameTerminalRequest
) -> TerminalInfoResponse:
    """
    Rename a terminal.

    Args:
        project_name: Name of the project
        terminal_id: ID of the terminal to rename
        request: Request body with new name

    Returns:
        The updated terminal info
    """
    if not validate_project_name(project_name):
        raise HTTPException(status_code=400, detail="Invalid project name")

    if not validate_terminal_id(terminal_id):
        raise HTTPException(status_code=400, detail="Invalid terminal ID")

    project_dir = _get_project_path(project_name)
    if not project_dir:
        raise HTTPException(status_code=404, detail="Project not found")

    if not rename_terminal(project_name, terminal_id, request.name):
        raise HTTPException(status_code=404, detail="Terminal not found")

    info = get_terminal_info(project_name, terminal_id)
    if not info:
        raise HTTPException(status_code=404, detail="Terminal not found")

    return TerminalInfoResponse(id=info.id, name=info.name, created_at=info.created_at)


@router.delete("/{project_name}/{terminal_id}")
async def delete_project_terminal(project_name: str, terminal_id: str) -> dict:
    """
    Delete a terminal and stop its session.

    Args:
        project_name: Name of the project
        terminal_id: ID of the terminal to delete

    Returns:
        Success message
    """
    if not validate_project_name(project_name):
        raise HTTPException(status_code=400, detail="Invalid project name")

    if not validate_terminal_id(terminal_id):
        raise HTTPException(status_code=400, detail="Invalid terminal ID")

    project_dir = _get_project_path(project_name)
    if not project_dir:
        raise HTTPException(status_code=404, detail="Project not found")

    # Stop the session if it's running
    await stop_terminal_session(project_name, terminal_id)

    # Delete the terminal metadata
    if not delete_terminal(project_name, terminal_id):
        raise HTTPException(status_code=404, detail="Terminal not found")

    return {"message": "Terminal deleted"}


# WebSocket Endpoint


@router.websocket("/ws/{project_name}/{terminal_id}")
async def terminal_websocket(websocket: WebSocket, project_name: str, terminal_id: str) -> None:
    """
    WebSocket endpoint for interactive terminal I/O.

    Message protocol:

    Client -> Server:
    - {"type": "input", "data": "<base64-encoded-bytes>"} - Keyboard input
    - {"type": "resize", "cols": 80, "rows": 24} - Terminal resize
    - {"type": "ping"} - Keep-alive ping

    Server -> Client:
    - {"type": "output", "data": "<base64-encoded-bytes>"} - PTY output
    - {"type": "exit", "code": 0} - Shell process exited
    - {"type": "pong"} - Keep-alive response
    - {"type": "error", "message": "..."} - Error message
    """
    # Always accept WebSocket first to avoid opaque 403 errors
    await websocket.accept()

    # Validate project name
    try:
        project_name = validate_project_name(project_name)
    except Exception:
        await websocket.send_json({"type": "error", "message": "Invalid project name"})
        await websocket.close(
            code=TerminalCloseCode.INVALID_PROJECT_NAME, reason="Invalid project name"
        )
        return

    # Validate terminal ID
    if not validate_terminal_id(terminal_id):
        await websocket.send_json({"type": "error", "message": "Invalid terminal ID"})
        await websocket.close(
            code=TerminalCloseCode.INVALID_PROJECT_NAME, reason="Invalid terminal ID"
        )
        return

    # Look up project directory from registry
    project_dir = _get_project_path(project_name)
    if not project_dir:
        await websocket.send_json({"type": "error", "message": "Project not found in registry"})
        await websocket.close(
            code=TerminalCloseCode.PROJECT_NOT_FOUND,
            reason="Project not found in registry",
        )
        return

    if not project_dir.exists():
        await websocket.send_json({"type": "error", "message": "Project directory not found"})
        await websocket.close(
            code=TerminalCloseCode.PROJECT_NOT_FOUND,
            reason="Project directory not found",
        )
        return

    # Verify terminal exists in metadata
    terminal_info = get_terminal_info(project_name, terminal_id)
    if not terminal_info:
        await websocket.send_json({"type": "error", "message": "Terminal not found"})
        await websocket.close(
            code=TerminalCloseCode.PROJECT_NOT_FOUND,
            reason="Terminal not found",
        )
        return

    # Get or create terminal session for this project/terminal
    session = get_terminal_session(project_name, project_dir, terminal_id)

    # Queue for output data to send to client
    output_queue: asyncio.Queue[bytes] = asyncio.Queue()

    # Callback to receive terminal output and queue it for sending
    def on_output(data: bytes) -> None:
        """Queue terminal output for async sending to WebSocket."""
        try:
            output_queue.put_nowait(data)
        except asyncio.QueueFull:
            logger.warning(f"Output queue full for {project_name}, dropping data")

    # Register the output callback
    session.add_output_callback(on_output)

    # Track if we need to wait for initial resize before starting
    # This ensures the PTY is created with correct dimensions from the start
    needs_initial_resize = not session.is_active

    # Task to send queued output to WebSocket
    async def send_output_task() -> None:
        """Continuously send queued output to the WebSocket client."""
        try:
            while True:
                # Wait for output data
                data = await output_queue.get()

                # Encode as base64 and send
                encoded = base64.b64encode(data).decode("ascii")
                await websocket.send_json({"type": "output", "data": encoded})

        except asyncio.CancelledError:
            raise
        except WebSocketDisconnect:
            raise
        except Exception as e:
            logger.warning(f"Error sending output for {project_name}: {e}")
            raise

    # Task to monitor if the terminal session exits
    async def monitor_exit_task() -> None:
        """Monitor the terminal session and notify client on exit."""
        try:
            # Wait for session to become active first (deferred start)
            while not session.is_active:
                await asyncio.sleep(0.1)

            # Now monitor until it becomes inactive
            while session.is_active:
                await asyncio.sleep(0.5)

            # Session ended - send exit message
            # Note: We don't have access to actual exit code from PTY
            await websocket.send_json({"type": "exit", "code": 0})

        except asyncio.CancelledError:
            raise
        except WebSocketDisconnect:
            raise
        except Exception as e:
            logger.warning(f"Error in exit monitor for {project_name}: {e}")

    # Start background tasks
    output_task = asyncio.create_task(send_output_task())
    exit_task = asyncio.create_task(monitor_exit_task())

    try:
        while True:
            try:
                # Receive message from client
                data = await websocket.receive_text()
                message = json.loads(data)
                msg_type = message.get("type")

                if msg_type == "ping":
                    await websocket.send_json({"type": "pong"})

                elif msg_type == "input":
                    # Only allow input after terminal is started
                    if not session.is_active:
                        await websocket.send_json(
                            {"type": "error", "message": "Terminal not ready - send resize first"}
                        )
                        continue

                    # Decode base64 input and write to PTY
                    encoded_data = message.get("data", "")
                    # Add size limit to prevent DoS
                    if len(encoded_data) > 65536:  # 64KB limit for base64 encoded data
                        await websocket.send_json({"type": "error", "message": "Input too large"})
                        continue
                    if encoded_data:
                        try:
                            decoded = base64.b64decode(encoded_data)
                        except (ValueError, TypeError) as e:
                            logger.warning(f"Failed to decode base64 input: {e}")
                            await websocket.send_json(
                                {"type": "error", "message": "Invalid base64 data"}
                            )
                            continue

                        try:
                            session.write(decoded)
                        except Exception as e:
                            logger.warning(f"Failed to write to terminal: {e}")
                            await websocket.send_json(
                                {"type": "error", "message": "Failed to write to terminal"}
                            )

                elif msg_type == "resize":
                    # Resize the terminal
                    cols = message.get("cols", 80)
                    rows = message.get("rows", 24)

                    # Validate dimensions
                    if isinstance(cols, int) and isinstance(rows, int):
                        cols = max(10, min(500, cols))
                        rows = max(5, min(200, rows))

                        # If this is the first resize and session not started, start with these dimensions
                        # This ensures the PTY is created with correct size from the beginning
                        if needs_initial_resize and not session.is_active:
                            started = await session.start(cols=cols, rows=rows)
                            if not started:
                                session.remove_output_callback(on_output)
                                try:
                                    await websocket.send_json(
                                        {"type": "error", "message": "Failed to start terminal session"}
                                    )
                                except Exception:
                                    pass
                                await websocket.close(
                                    code=TerminalCloseCode.FAILED_TO_START, reason="Failed to start terminal"
                                )
                                return
                            # Mark that we no longer need initial resize
                            needs_initial_resize = False
                        else:
                            session.resize(cols, rows)
                    else:
                        await websocket.send_json({"type": "error", "message": "Invalid resize dimensions"})

                else:
                    await websocket.send_json({"type": "error", "message": f"Unknown message type: {msg_type}"})

            except json.JSONDecodeError:
                await websocket.send_json({"type": "error", "message": "Invalid JSON"})

    except WebSocketDisconnect:
        logger.info(f"Terminal WebSocket disconnected for {project_name}/{terminal_id}")

    except Exception as e:
        logger.exception(f"Terminal WebSocket error for {project_name}/{terminal_id}")
        try:
            await websocket.send_json({"type": "error", "message": f"Server error: {str(e)}"})
        except Exception:
            pass

    finally:
        # Cancel background tasks
        output_task.cancel()
        exit_task.cancel()

        try:
            await output_task
        except asyncio.CancelledError:
            pass

        try:
            await exit_task
        except asyncio.CancelledError:
            pass

        # Remove the output callback
        session.remove_output_callback(on_output)

        # Only stop session if no other clients are connected
        with session._callbacks_lock:
            remaining_callbacks = len(session._output_callbacks)

        if remaining_callbacks == 0:
            await session.stop()
            logger.info(f"Terminal session stopped for {project_name}/{terminal_id} (last client disconnected)")
        else:
            logger.info(
                f"Client disconnected from {project_name}/{terminal_id}, {remaining_callbacks} clients remaining"
            )