headquarter/apps/api/src/api/system/terminal.py

"""WebSocket terminal endpoint for tool instances."""

import asyncio
import json
import logging
import uuid
from contextlib import suppress

from fastapi import APIRouter, Depends, HTTPException, WebSocket, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from starlette.websockets import WebSocketDisconnect

from src.auth.dependencies import get_current_user_id, get_db_session
from src.models import TerminalSessionModel
from src.models import ToolInstance
from src.models import ToolType
from src.services.terminal.terminal_manager import MaxSessionsExceededError, terminal_manager

router = APIRouter()
logger = logging.getLogger(__name__)


class SessionRef:
    """Mutable reference to a terminal session, allowing updates during reset."""

    def __init__(self, session, slot_session_id: str | None = None):
        self.session = session
        self.slot_session_id = slot_session_id or session.session_id


@router.websocket(
    "/ws/tool-instances/{instance_id}/terminal",
)
async def terminal_websocket_default(
    websocket: WebSocket,
    instance_id: str,
    db_session: AsyncSession = Depends(get_db_session),
) -> None:
    """WebSocket endpoint for terminal access (default session alias).

    Backward-compatible route that maps to the default session.
    """
    await _handle_terminal_websocket(websocket, instance_id, None, db_session)


@router.websocket(
    "/ws/tool-instances/{instance_id}/terminal/{session_id}",
)
async def terminal_websocket_specific(
    websocket: WebSocket,
    instance_id: str,
    session_id: str,
    db_session: AsyncSession = Depends(get_db_session),
) -> None:
    """WebSocket endpoint for a specific terminal session."""
    await _handle_terminal_websocket(websocket, instance_id, session_id, db_session)


async def _handle_terminal_websocket(
    websocket: WebSocket,
    instance_id: str,
    target_session_id: str | None,
    db_session: AsyncSession,
) -> None:
    """Shared WebSocket handler for terminal sessions.

    Args:
        websocket: The WebSocket connection.
        instance_id: UUID string of the tool instance.
        target_session_id: Specific session ID (slot key). None means default session.
        db_session: Database session.
    """
    logger.debug(
        "Terminal WebSocket connection attempt for instance %s (session=%s)",
        instance_id,
        target_session_id or "default",
    )
    await websocket.accept()
    logger.debug("Terminal WebSocket accepted for instance %s", instance_id)

    try:
        # Parse instance_id
        instance_uuid = uuid.UUID(instance_id)
    except ValueError:
        logger.error("Invalid instance ID: %s", instance_id)
        await websocket.close(code=4001, reason="Invalid instance ID")
        return

    # Authenticate user from session cookie
    user_id = await _get_user_from_websocket(websocket, db_session)
    if user_id is None:
        logger.warning(
            "Unauthorized terminal access attempt for instance %s", instance_id
        )
        await websocket.close(code=4003, reason="Unauthorized")
        return

    # Get instance and verify ownership
    instance = await db_session.get(ToolInstance, instance_uuid)
    if instance is None:
        logger.warning("Instance %s not found", instance_id)
        await websocket.close(code=4004, reason="Instance not found")
        return

    if instance.owner_id != user_id:
        logger.warning(
            "Forbidden terminal access for instance %s by user %s",
            instance_id,
            user_id,
        )
        await websocket.close(code=4003, reason="Forbidden")
        return

    if instance.status != "running" or not instance.container_id:
        logger.warning(
            "Instance %s not running (status=%s, container_id=%s)",
            instance_id,
            instance.status,
            instance.container_id,
        )
        await websocket.close(code=4004, reason="Instance not running")
        return

    logger.debug("Terminal auth passed for instance %s, user %s", instance_id, user_id)

    # Verify the container actually exists (may have been removed/recreated)
    from src.services.docker import get_container_status

    container_status = get_container_status(instance.container_id)
    if container_status["status"] == "not_found":
        logger.error(
            "Container %s for instance %s not found (may have been removed)",
            instance.container_id,
            instance_id,
        )
        await websocket.close(
            code=4004, reason="Container not found — restart the tool instance"
        )
        return

    # Fetch tool type to get startup_command
    tool_type = await db_session.get(ToolType, instance.tool_type_id)
    startup_command = tool_type.startup_command if tool_type else None
    if startup_command:
        logger.debug(
            "Using startup command for instance %s: %s",
            instance_id,
            startup_command,
        )

    session = None

    # Get or create terminal session
    try:
        if target_session_id is None:
            # Default session alias
            session = await terminal_manager.get_or_create_session(
                instance_uuid,
                instance.container_id,
                startup_command=startup_command,
            )
            slot_session_id = "default"
        else:
            # Specific session
            session = terminal_manager.get_session(
                instance_id,
                target_session_id,
            )
            if session is None:
                # Session not in memory — may have been lost on server restart.
                # Try to restore from the DB row.
                db_row = await db_session.get(
                    TerminalSessionModel, uuid.UUID(target_session_id)
                )
                if (
                    db_row is not None
                    and db_row.instance_id == instance_uuid
                    and db_row.status != "closed"
                ):
                    logger.info(
                        "Restoring terminal session %s for instance %s from DB",
                        target_session_id,
                        instance_id,
                    )
                    session = await terminal_manager.create_session(
                        instance_uuid,
                        instance.container_id,
                        startup_command=startup_command,
                        name=db_row.name,
                        session_id=target_session_id,
                    )
                else:
                    logger.warning(
                        "Session %s not found for instance %s",
                        target_session_id,
                        instance_id,
                    )
                    await websocket.close(code=4004, reason="Session not found")
                    return
            # Determine slot key for reset scoping
            key = terminal_manager._find_key_by_internal_id(
                instance_id, session.session_id
            )
            slot_session_id = key[1] if key else target_session_id

        logger.debug(
            "Terminal session ready for instance %s (session_id=%s, slot=%s)",
            instance_id,
            session.session_id,
            slot_session_id,
        )

        # Attach WebSocket to session
        await terminal_manager.attach_websocket(session, websocket)
        logger.debug("WebSocket attached to session for instance %s", instance_id)

        # Send connected status
        await websocket.send_json({"type": "status", "status": "connected"})
        logger.debug("Sent connected status for instance %s", instance_id)

        # Use mutable session reference so loops can survive reset
        session_ref = SessionRef(session, slot_session_id)

        # Start write loop and heartbeat (read is now event-driven in TerminalSession)
        write_task = asyncio.create_task(
            _write_loop(session_ref, websocket, instance_id)
        )
        heartbeat_task = asyncio.create_task(_heartbeat_loop(websocket))
        logger.debug("Started terminal loops for instance %s", instance_id)

        # Wait for either task to complete (indicating disconnect or error)
        done, pending = await asyncio.wait(
            [write_task, heartbeat_task],
            return_when=asyncio.FIRST_COMPLETED,
        )

        logger.debug(
            "Terminal loop completed for instance %s, done=%s",
            instance_id,
            len(done),
        )

        # Cancel remaining tasks
        for task in pending:
            task.cancel()

    except WebSocketDisconnect:
        logger.debug("WebSocket disconnected for instance %s", instance_id)
    except Exception as exc:
        logger.error(
            "Terminal session error for instance %s: %s",
            instance_id,
            str(exc),
            exc_info=True,
        )
        with suppress(Exception):
            await websocket.close(code=4000, reason=f"Error: {exc}")
    finally:
        # Detach WebSocket, don't kill session
        with suppress(Exception):
            if session is not None:
                await terminal_manager.detach_websocket(session, websocket)
                logger.debug(
                    "WebSocket detached from session for instance %s", instance_id
                )


async def _write_loop(session_ref: SessionRef, websocket, instance_id: str) -> None:
    """Read input from WebSocket and send to container."""
    try:
        while True:
            session = session_ref.session
            if not session.is_alive() or session._closed:
                await asyncio.sleep(0.1)
                continue
            message = await websocket.receive()
            if message["type"] == "websocket.receive":
                if "bytes" in message:
                    await session.write_input(message["bytes"])
                elif "text" in message:
                    text = message["text"]
                    if text.startswith("{"):
                        # Control message (JSON)
                        try:
                            ctrl = json.loads(text)
                            msg_type = ctrl.get("type")

                            if msg_type == "resize":
                                cols = ctrl.get("cols", 80)
                                rows = ctrl.get("rows", 24)
                                logger.debug(
                                    "Received resize message for instance %s: %sx%s",
                                    instance_id,
                                    cols,
                                    rows,
                                )
                                await session.resize(cols, rows)
                            elif msg_type == "ack":
                                char_count = ctrl.get("chars", 0)
                                if char_count > 0:
                                    session.acknowledge_data(char_count)
                            elif msg_type == "reset":
                                # Reset terminal session (scoped to current slot)
                                logger.debug(
                                    "Resetting terminal session for instance %s (slot=%s)",
                                    session.instance_id,
                                    session_ref.slot_session_id,
                                )
                                await websocket.send_json(
                                    {"type": "status", "status": "resetting"}
                                )

                                # Reset the session scoped to its slot
                                new_session = await terminal_manager.reset_session(
                                    session.instance_id,
                                    session.container_id,
                                    startup_command=session.startup_command,
                                    session_id=session_ref.slot_session_id,
                                    name=session.name,
                                )

                                # Update the mutable session reference
                                session_ref.session = new_session

                                # Attach to new session
                                await terminal_manager.attach_websocket(
                                    new_session, websocket
                                )
                                await websocket.send_json(
                                    {"type": "status", "status": "connected"}
                                )

                                # Continue the loop with the new session
                                continue

                        except json.JSONDecodeError:
                            # Not a valid JSON control message, treat as regular input
                            await session.write_input(text.encode("utf-8"))
                    else:
                        await session.write_input(text.encode("utf-8"))
            elif message["type"] == "websocket.disconnect":
                break
    except Exception:
        pass


async def _heartbeat_loop(websocket: WebSocket) -> None:
    """Send periodic ping messages to detect disconnections."""
    try:
        while True:
            await asyncio.sleep(30)  # Ping every 30 seconds
            try:
                await websocket.send_json({"type": "ping"})
            except Exception:
                # WebSocket is closed or broken
                break
    except Exception:
        pass


async def _get_terminal_instance(
    instance_id: uuid.UUID,
    user_id: uuid.UUID,
    db_session: AsyncSession,
) -> ToolInstance:
    """Fetch instance and validate auth, ownership, and running status.

    Args:
        instance_id: UUID of the tool instance.
        user_id: ID of the authenticated user.
        db_session: Database session.

    Returns:
        The validated ToolInstance.

    Raises:
        HTTPException: If instance not found, not owned, or not running.
    """
    instance = await db_session.get(ToolInstance, instance_id)
    if instance is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Instance not found"
        )

    if instance.owner_id != user_id:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Not authorized to access this instance",
        )

    if instance.status != "running" or not instance.container_id:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail="Instance is not running"
        )

    return instance


@router.get(
    "/instances/{instance_id}/terminal/sessions",
    summary="List terminal sessions",
    description="List terminal sessions for a tool instance with live WebSocket state.",
)
async def list_terminal_sessions(
    instance_id: uuid.UUID,
    user_id: uuid.UUID = Depends(get_current_user_id),
    db_session: AsyncSession = Depends(get_db_session),
) -> dict:
    """List terminal sessions for an instance.

    Args:
        instance_id: UUID of the tool instance.
        user_id: ID of the authenticated user.
        db_session: Database session.

    Returns:
        Dictionary with sessions list.
    """
    await _get_terminal_instance(instance_id, user_id, db_session)

    # Query active DB rows for this instance
    result = await db_session.execute(
        select(TerminalSessionModel)
        .where(TerminalSessionModel.instance_id == instance_id)
        .where(TerminalSessionModel.status != "closed")
        .order_by(TerminalSessionModel.created_at.asc())
    )
    db_rows = result.scalars().all()

    # Build response with live has_websockets flag.
    # Include DB rows even without in-memory counterparts (e.g. after
    # server restart) so the frontend can display tabs and reconnect.
    sessions = []
    for row in db_rows:
        live_session = terminal_manager.get_session(str(instance_id), str(row.id))
        sessions.append(
            {
                "id": str(row.id),
                "name": row.name,
                "status": row.status,
                "has_websockets": live_session.has_websockets()
                if live_session
                else False,
                "created_at": row.created_at.isoformat() if row.created_at else None,
                "last_activity_at": row.last_activity_at.isoformat()
                if row.last_activity_at
                else None,
            }
        )

    return {"sessions": sessions}


@router.post(
    "/instances/{instance_id}/terminal/sessions",
    summary="Create terminal session",
    description="Create a new terminal session for a running tool instance.",
    status_code=status.HTTP_201_CREATED,
)
async def create_terminal_session(
    instance_id: uuid.UUID,
    data: dict,
    user_id: uuid.UUID = Depends(get_current_user_id),
    db_session: AsyncSession = Depends(get_db_session),
) -> dict:
    """Create a new terminal session.

    Args:
        instance_id: UUID of the tool instance.
        data: Request body with optional name.
        user_id: ID of the authenticated user.
        db_session: Database session.

    Returns:
        Dictionary with new session details.

    Raises:
        HTTPException: 409 if max sessions reached.
    """
    instance = await _get_terminal_instance(instance_id, user_id, db_session)
    assert instance.container_id is not None

    # Fetch tool type to get startup_command
    tool_type = await db_session.get(ToolType, instance.tool_type_id)
    startup_command = tool_type.startup_command if tool_type else None

    name = data.get("name")

    try:
        session = await terminal_manager.create_session(
            instance_id,
            instance.container_id,
            startup_command=startup_command,
            name=name,
        )
    except MaxSessionsExceededError:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail="Maximum of 5 terminal sessions reached for this instance",
        ) from None

    return {
        "id": session.session_id,
        "name": session.name,
        "status": session.status,
        "created_at": session.last_activity,
    }


@router.delete(
    "/instances/{instance_id}/terminal/sessions/{session_id}",
    summary="Close terminal session",
    description="Close a specific terminal session.",
)
async def close_terminal_session(
    instance_id: uuid.UUID,
    session_id: str,
    user_id: uuid.UUID = Depends(get_current_user_id),
    db_session: AsyncSession = Depends(get_db_session),
) -> dict:
    """Close a terminal session.

    Args:
        instance_id: UUID of the tool instance.
        session_id: ID of the session to close.
        user_id: ID of the authenticated user.
        db_session: Database session.

    Returns:
        Dictionary with closure status.
    """
    await _get_terminal_instance(instance_id, user_id, db_session)

    # Find the session by internal ID to determine its slot key
    key = terminal_manager._find_key_by_internal_id(str(instance_id), session_id)
    if (
        key is None
        and terminal_manager.get_session(str(instance_id), session_id) is not None
    ):
        key = (str(instance_id), session_id)

    if key is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Session not found"
        )

    await terminal_manager.close_session(key[0], key[1])

    return {"status": "closed", "session_id": session_id}


@router.post(
    "/instances/{instance_id}/terminal/sessions/{session_id}/reset",
    summary="Reset terminal session",
    description="Reset a specific terminal session, killing the current shell and starting fresh.",
)
async def reset_specific_terminal_session(
    instance_id: uuid.UUID,
    session_id: str,
    user_id: uuid.UUID = Depends(get_current_user_id),
    db_session: AsyncSession = Depends(get_db_session),
) -> dict:
    """Reset a specific terminal session.

    Args:
        instance_id: UUID of the tool instance.
        session_id: ID of the session to reset.
        user_id: ID of the authenticated user.
        db_session: Database session.

    Returns:
        Dictionary with reset session details.
    """
    instance = await _get_terminal_instance(instance_id, user_id, db_session)
    assert instance.container_id is not None

    # Determine slot key for reset
    key = terminal_manager._find_key_by_internal_id(str(instance_id), session_id)
    if (
        key is None
        and terminal_manager.get_session(str(instance_id), session_id) is not None
    ):
        key = (str(instance_id), session_id)

    if key is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Session not found"
        )

    # Fetch tool type to get startup_command
    tool_type = await db_session.get(ToolType, instance.tool_type_id)
    startup_command = tool_type.startup_command if tool_type else None

    # Preserve name if possible
    live_session = terminal_manager.get_session(str(instance_id), session_id)
    name = live_session.name if live_session else None

    new_session = await terminal_manager.reset_session(
        instance_id,
        instance.container_id,
        startup_command=startup_command,
        session_id=key[1],
        name=name,
    )

    return {
        "id": new_session.session_id,
        "name": new_session.name,
        "status": new_session.status,
    }


@router.post(
    "/instances/{instance_id}/terminal/sessions/{session_id}/rename",
    summary="Rename terminal session",
    description="Rename a specific terminal session.",
)
async def rename_terminal_session(
    instance_id: uuid.UUID,
    session_id: str,
    data: dict,
    user_id: uuid.UUID = Depends(get_current_user_id),
    db_session: AsyncSession = Depends(get_db_session),
) -> dict:
    """Rename a terminal session.

    Args:
        instance_id: UUID of the tool instance.
        session_id: ID of the session to rename.
        data: Request body with new name.
        user_id: ID of the authenticated user.
        db_session: Database session.

    Returns:
        Dictionary with updated session details.
    """
    await _get_terminal_instance(instance_id, user_id, db_session)

    new_name = data.get("name")
    if not new_name or not isinstance(new_name, str):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail="Name is required"
        )

    # Update in-memory session name if live
    live_session = terminal_manager.get_session(str(instance_id), session_id)
    if live_session:
        live_session.name = new_name

    # Update DB row
    db_row = await db_session.get(TerminalSessionModel, uuid.UUID(session_id))
    if db_row is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Session not found"
        )

    db_row.name = new_name
    await db_session.commit()

    return {"id": str(db_row.id), "name": new_name}


@router.post(
    "/instances/{instance_id}/terminal/reset",
    summary="Reset terminal session (legacy alias)",
    description="Reset the default terminal session for a tool instance. Preserved for backward compatibility.",
)
async def reset_terminal_session(
    instance_id: uuid.UUID,
    user_id: uuid.UUID = Depends(get_current_user_id),
    db_session: AsyncSession = Depends(get_db_session),
) -> dict:
    """Reset the default terminal session for an instance (legacy alias).

    Args:
        instance_id: UUID of the tool instance.
        user_id: ID of the authenticated user.
        db_session: Database session.

    Returns:
        Dictionary with status message.
    """
    instance = await _get_terminal_instance(instance_id, user_id, db_session)
    assert instance.container_id is not None

    # Fetch tool type to get startup_command
    tool_type = await db_session.get(ToolType, instance.tool_type_id)
    startup_command = tool_type.startup_command if tool_type else None

    try:
        # Reset the default session
        new_session = await terminal_manager.reset_session(
            instance_id,
            instance.container_id,
            startup_command=startup_command,
        )

        logger.info(
            "Terminal session reset for instance %s (new session_id=%s)",
            instance_id,
            new_session.session_id,
        )

        return {
            "status": "success",
            "message": "Terminal session reset successfully",
            "instance_id": str(instance_id),
            "session_id": new_session.session_id,
        }
    except Exception as exc:
        logger.error(
            "Failed to reset terminal session for instance %s: %s",
            instance_id,
            str(exc),
            exc_info=True,
        )
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to reset terminal session: {exc}",
        ) from exc


async def _get_user_from_websocket(
    websocket: WebSocket,
    db_session: AsyncSession,
) -> uuid.UUID | None:
    """Extract and validate user ID from session cookie in WebSocket.

    Args:
        websocket: The WebSocket connection.
        db_session: Database session.

    Returns:
        The user's UUID if authenticated, None otherwise.
    """
    from src.auth.session import decode_session_cookie
    from src.config import Settings

    session_cookie = websocket.cookies.get("session")
    if not session_cookie:
        return None

    settings = Settings()
    try:
        payload = decode_session_cookie(settings=settings, cookie_value=session_cookie)
        return uuid.UUID(str(payload["user_id"]))
    except (ValueError, KeyError):
        return None