mybeacon-backend/backend/app/api/v1/superadmin/tunnels.py

"""
Superadmin tunnel management API endpoints.
"""

from typing import Annotated, Optional

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.ext.asyncio import AsyncSession
from pydantic import BaseModel
from sqlalchemy import select

from app.api.deps import get_current_superadmin
from app.core.database import get_db
from app.models.device import Device
from app.models.user import User
from app.services.tunnel_service import tunnel_service


router = APIRouter(prefix="/tunnels", tags=["superadmin-tunnels"])


class TunnelEnableResponse(BaseModel):
    """Response when enabling tunnel"""
    session_uuid: str
    device_id: str
    tunnel_type: str
    status: str


class TunnelStatusResponse(BaseModel):
    """Tunnel session status response"""
    session_uuid: str
    status: str  # "waiting" | "ready" | "failed"
    device_tunnel_port: Optional[int] = None
    ttyd_port: Optional[int] = None
    tunnel_url: Optional[str] = None


@router.post("/devices/{device_id}/{tunnel_type}")
async def enable_tunnel(
    device_id: int,
    tunnel_type: str,
    db: Annotated[AsyncSession, Depends(get_db)],
    current_user: Annotated[User, Depends(get_current_superadmin)]
) -> TunnelEnableResponse:
    """
    Enable SSH or Dashboard tunnel for device.
    Creates session and triggers device to connect tunnel.

    Args:
        device_id: Device numeric ID
        tunnel_type: "ssh" or "dashboard"

    Returns:
        session_uuid: UUID for polling status
    """
    if tunnel_type not in ["ssh", "dashboard"]:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="tunnel_type must be 'ssh' or 'dashboard'"
        )

    # Get device
    result = await db.execute(
        select(Device).where(Device.id == device_id)
    )
    device = result.scalar_one_or_none()

    if not device:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Device not found"
        )

    # Create tunnel session
    session = tunnel_service.create_session(
        device_id=device.mac_address,
        admin_user=current_user.email,
        tunnel_type=tunnel_type
    )

    # Update device config to enable tunnel
    if device.config is None:
        device.config = {}

    tunnel_key = f"{tunnel_type}_tunnel"
    if tunnel_key not in device.config:
        device.config[tunnel_key] = {}

    device.config[tunnel_key]["enabled"] = True

    # Mark as modified (SQLAlchemy JSON field)
    from sqlalchemy.orm import attributes
    attributes.flag_modified(device, "config")

    await db.commit()

    return TunnelEnableResponse(
        session_uuid=session.uuid,
        device_id=device.mac_address,
        tunnel_type=tunnel_type,
        status="waiting"
    )


@router.get("/sessions/{session_uuid}/status")
async def get_tunnel_status(
    session_uuid: str,
    current_user: Annotated[User, Depends(get_current_superadmin)]
) -> TunnelStatusResponse:
    """
    Poll tunnel session status.

    Returns:
        - waiting: Device not yet connected
        - ready: Tunnel connected, ttyd spawned
        - failed: Session expired or failed
    """
    session = tunnel_service.get_session(session_uuid)

    if not session:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Session not found or expired"
        )

    # Build response
    response = TunnelStatusResponse(
        session_uuid=session.uuid,
        status=session.status,
        device_tunnel_port=session.device_tunnel_port,
        ttyd_port=session.ttyd_port
    )

    # If ready, spawn ttyd if not already spawned
    if session.status == "ready" and session.device_tunnel_port:
        if not session.ttyd_port:
            try:
                # Spawn ttyd
                ttyd_port = tunnel_service.spawn_ttyd(
                    session_uuid=session.uuid,
                    device_tunnel_port=session.device_tunnel_port
                )
                response.ttyd_port = ttyd_port
                response.tunnel_url = f"/admin/{session.tunnel_type}/{session.uuid}"
            except Exception as e:
                print(f"[tunnel] Failed to spawn ttyd: {e}")
                session.status = "failed"
                response.status = "failed"
        else:
            response.tunnel_url = f"/admin/{session.tunnel_type}/{session.uuid}"

    return response


@router.post("/sessions/{session_uuid}/heartbeat")
async def session_heartbeat(
    session_uuid: str,
    current_user: Annotated[User, Depends(get_current_superadmin)]
):
    """
    Browser sends heartbeat every 30 seconds to keep session alive.
    """
    success = tunnel_service.update_heartbeat(session_uuid)

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

    return {"success": True}