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

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

from typing import Optional

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

from app.core.database import get_db
from app.core.permissions import require_permission
from app.models.device import Device
from app.models.user import User
from app.services.tunnel_service import tunnel_service
from sqlalchemy import select


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}")
@require_permission("devices", "manage")
async def enable_tunnel(
    device_id: int,
    tunnel_type: str,
    current_user: User = Depends(require_permission("devices", "manage")),
    db: AsyncSession = Depends(get_db)
) -> 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: User = Depends(require_permission("devices", "view"))
) -> 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: User = Depends(require_permission("devices", "view"))
):
    """
    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}