SmartUp/backend/app/routers/browser_sessions.py

"""Remote browser session API."""
from __future__ import annotations

import asyncio
import hashlib
import json
import logging
from typing import Any, Literal, Optional

from fastapi import APIRouter, Depends, HTTPException, Query, WebSocket, WebSocketDisconnect
from fastapi.responses import Response
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from app.database import get_db
from app.models.custom_page import CustomPage
from app.services.browser_session_service import (
    BrowserDependencyError,
    BrowserSessionError,
    browser_sessions,
)
from app.utils.auth import decode_token, get_current_user, get_user_from_token_param
from app.database import SessionLocal
from app.models.admin_user import AdminUser

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/browser-sessions", tags=["browser-sessions"])


class BrowserSessionCreate(BaseModel):
    custom_page_id: int
    width: int = Field(default=1280)
    height: int = Field(default=720)


class BrowserSessionResponse(BaseModel):
    id: str
    custom_page_id: int
    url: str
    title: str


class BrowserEvent(BaseModel):
    type: Literal["click", "dblclick", "mousemove", "mousedown", "mouseup", "type", "key", "scroll", "reload", "back", "forward", "resize"]
    x: Optional[float] = None
    y: Optional[float] = None
    button: Optional[Literal["left", "right", "middle"]] = "left"
    text: Optional[str] = None
    key: Optional[str] = None
    delta_x: Optional[float] = 0
    delta_y: Optional[float] = 0
    width: Optional[int] = None
    height: Optional[int] = None


def _error_from_browser(exc: Exception) -> HTTPException:
    if isinstance(exc, BrowserDependencyError):
        return HTTPException(503, str(exc))
    if isinstance(exc, BrowserSessionError):
        return HTTPException(409, str(exc))
    if isinstance(exc, KeyError):
        return HTTPException(404, "browser session not found")
    if isinstance(exc, ValueError):
        return HTTPException(400, str(exc))
    return HTTPException(502, f"Browser error: {exc}")


@router.post("", response_model=BrowserSessionResponse, status_code=201)
async def create_session(
    body: BrowserSessionCreate,
    db: Session = Depends(get_db),
    _=Depends(get_current_user),
):
    page = db.query(CustomPage).filter(CustomPage.id == body.custom_page_id).first()
    if not page or not page.enabled:
        raise HTTPException(404, "page not found")
    if page.access_mode != "remote_browser":
        raise HTTPException(400, "custom page is not configured for remote browser mode")
    login_config = {
        "enabled": page.login_autofill_enabled,
        "username": page.login_username,
        "password": page.login_password,
        "username_selector": page.login_username_selector,
        "password_selector": page.login_password_selector,
        "submit_selector": page.login_submit_selector,
    }
    try:
        session = await browser_sessions.create(page.id, page.url, body.width, body.height, login_config)
        return await browser_sessions.state(session.id)
    except Exception as exc:
        raise _error_from_browser(exc)


@router.get("/{session_id}", response_model=BrowserSessionResponse)
async def get_session(session_id: str, _=Depends(get_current_user)):
    try:
        return await browser_sessions.state(session_id)
    except Exception as exc:
        raise _error_from_browser(exc)


@router.get("/{session_id}/screenshot")
async def session_screenshot(session_id: str, _=Depends(get_user_from_token_param)):
    try:
        image = await browser_sessions.screenshot(session_id)
    except Exception as exc:
        raise _error_from_browser(exc)
    return Response(content=image, media_type="image/jpeg", headers={"Cache-Control": "no-store"})


@router.post("/{session_id}/events", response_model=BrowserSessionResponse)
async def send_event(session_id: str, body: BrowserEvent, _=Depends(get_current_user)):
    try:
        payload: dict[str, Any] = body.model_dump(exclude_none=True)
        event_type = payload.pop("type")
        return await browser_sessions.event(session_id, event_type, payload)
    except Exception as exc:
        raise _error_from_browser(exc)


@router.delete("/{session_id}", status_code=204)
async def close_session(session_id: str, _=Depends(get_current_user)):
    await browser_sessions.close(session_id)


# ——— WebSocket stream ———
# Frame interval & diff detection
_WS_MIN_INTERVAL = 0.05   # 50 ms floor (≈20 fps max)
_WS_IDLE_INTERVAL = 0.15  # 150 ms when nothing changed recently
_WS_ACTIVE_INTERVAL = 0.08  # 80 ms right after a user event


async def _ws_authenticate(token: Optional[str]) -> bool:
    """Validate JWT token for WebSocket connections."""
    if not token:
        return False
    email = decode_token(token)
    if not email:
        return False
    db = SessionLocal()
    try:
        user = db.query(AdminUser).filter(AdminUser.email == email).first()
        return user is not None
    finally:
        db.close()


@router.websocket("/{session_id}/ws")
async def session_ws(
    websocket: WebSocket,
    session_id: str,
    token: Optional[str] = Query(default=None),
):
    """WebSocket endpoint: pushes JPEG frames as binary, receives JSON event messages."""
    # Authenticate before accepting
    if not await _ws_authenticate(token):
        await websocket.close(code=4401)
        return

    await websocket.accept()

    # Track when a user event arrived so we can temporarily speed up
    last_event_at: float = 0.0
    last_frame_hash: str = ""

    # Task: receive events from client
    async def receive_loop():
        nonlocal last_event_at
        try:
            while True:
                raw = await websocket.receive_text()
                try:
                    msg = json.loads(raw)
                except json.JSONDecodeError:
                    continue
                msg_type = msg.get("type")
                if not msg_type:
                    continue
                payload: dict[str, Any] = {k: v for k, v in msg.items() if k != "type"}
                try:
                    await browser_sessions.event(session_id, msg_type, payload)
                    last_event_at = asyncio.get_event_loop().time()
                except Exception as exc:
                    logger.warning("ws event error: %s", exc)
                    try:
                        await websocket.send_json({"error": str(exc)})
                    except Exception:
                        pass
        except (WebSocketDisconnect, asyncio.CancelledError):
            pass
        except Exception as exc:
            logger.debug("ws receive_loop ended: %s", exc)

    # Task: push screenshots
    async def push_loop():
        nonlocal last_frame_hash
        try:
            while True:
                now = asyncio.get_event_loop().time()
                # Faster cadence right after a user interaction
                interval = _WS_ACTIVE_INTERVAL if (now - last_event_at) < 1.0 else _WS_IDLE_INTERVAL

                try:
                    frame = await browser_sessions.screenshot(session_id)
                except KeyError:
                    # Session gone
                    await websocket.send_json({"error": "session_not_found"})
                    break
                except Exception as exc:
                    logger.warning("ws screenshot error: %s", exc)
                    await asyncio.sleep(interval)
                    continue

                # Only push if content changed
                frame_hash = hashlib.md5(frame).hexdigest()
                if frame_hash != last_frame_hash:
                    last_frame_hash = frame_hash
                    try:
                        await websocket.send_bytes(frame)
                    except Exception:
                        break

                await asyncio.sleep(max(_WS_MIN_INTERVAL, interval))
        except (WebSocketDisconnect, asyncio.CancelledError):
            pass
        except Exception as exc:
            logger.debug("ws push_loop ended: %s", exc)

    # Send initial metadata so client knows session info
    try:
        state = await browser_sessions.state(session_id)
        await websocket.send_json({"type": "init", "session": state})
    except Exception as exc:
        await websocket.send_json({"error": f"session error: {exc}"})
        await websocket.close()
        return

    recv_task = asyncio.create_task(receive_loop())
    push_task = asyncio.create_task(push_loop())

    # Run until one side closes
    done, pending = await asyncio.wait(
        [recv_task, push_task],
        return_when=asyncio.FIRST_COMPLETED,
    )
    for t in pending:
        t.cancel()
        try:
            await t
        except asyncio.CancelledError:
            pass