AareDAQ/src/aare/daq/server.py

import asyncio
import hmac
import io
import os, time
import random
from contextlib import asynccontextmanager
from typing import AsyncGenerator
import json
import cv2
import urllib3
import uvicorn

from aare.common.coordinate import SmargonCoordinate, Coordinate, AerotechCoordinate
from aare.common.auth_models import BatonStatus, BatonRequestStatus
from aare.common.coordinate import SmargonCoordinate, Coordinate
from aare.common.error_codes import export_error_codes, export_error_codes_grouped
from aare.common.logger_config import setup_logger
from aare.common.models import SampleShortInfo, DAQStatusModel, BeamlineStateEnum, BeamlineSettingsModel, \
    SampleShortInfoList, SessionStatus, SampleCameraSettings, AutofocusSettings, TokenData, \
    CryojetSettingsModel, SimpleScanParameters, CrystalSize, FluorescenceSpectrumParameterModel, \
    FluorescenceSpectrumOutputModel, RecoveryActionRequest
from aare.common.raster_grid import RasterGridRequest, CompletedRasterGrid
from aare.common.rotation_scan import RotationScanRequest, CompletedRotationScan
from aare.common.sample_geometry import SampleGeometryModel
from fastapi import FastAPI, Depends, BackgroundTasks
from fastapi.concurrency import run_in_threadpool
from fastapi import HTTPException
from fastapi import status as api_status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from starlette.responses import StreamingResponse

from aare.daq import auth
from aare.common.beamline import mx_beamline
from aare.daq.config import BeamlineConfig
from aare.daq.daq import AareDAQ

from aare.daq.server_exception_handler import register_exception_handlers

from aare.common.exception_handler import (
    SampleException,
    UserRightsException,
    )

from aare.common.automation_queue_manager import WorkflowRedisManager
from aare.common.automation_workflow import STATE_REGISTRY, HANDLER_REGISTRY, SIMULATED_HANDLER_REGISTRY
from aare.daq.automation_runner import PersistentWorkflowRunner
from aare.daq.automation_api_router import router as workflow_router, set_workflow_dependencies

logger = setup_logger("aareDAQ")

# OAuth2 setup
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

# ── Per-worker state: populated inside the lifespan, after fork ──
bl = None
cfg = None
daq = None
workflow_redis_manager = None
workflow_runner = None

USE_SIMULATED_WORKFLOW = os.getenv("WORKFLOW_SIMULATION", "0") == "1"

_all_pgroups_cache: dict[str, tuple[list[str], float]] = {}
_ALL_PGROUPS_TTL_S = 60.0  # adjust TTL as needed

_face_detection_state: dict = {
    "seq": 0,
    "running": False,
    "samples": [],
    "height_fit": {},
    "area_fit": {},
}
_face_detection_state_lock = asyncio.Lock()

@asynccontextmanager
async def lifespan(application: FastAPI):
    """
    Runs once per worker process, AFTER fork() and BEFORE serving requests.
    All stateful / connection-opening initialisation belongs here so that
    each worker gets its own fresh Redis, BEC, EPICS, and TELL connections.
    """
    global bl, cfg, daq, workflow_redis_manager, workflow_runner

    await asyncio.sleep(random.uniform(0.5, 3.0))

    logger.info(f"Worker {os.getpid()} starting initialisation...")

    # ── Core objects (Redis, EPICS PVs, BEC, TELL, JFJoch, etc.) ──
    bl = mx_beamline()
    cfg = BeamlineConfig(bl)
    daq = AareDAQ(cfg, bl)

    # ── Workflow system ──
    # workflow_redis_manager = WorkflowRedisManager(
    #     client=cfg._BeamlineConfig__client,  # Reuse this worker's Redis connection
    #     beamline=bl.value,
    # )
    # workflow_runner = PersistentWorkflowRunner(
    #     redis_manager=workflow_redis_manager,
    #     registry=STATE_REGISTRY,
    #     handlers=SIMULATED_HANDLER_REGISTRY if USE_SIMULATED_WORKFLOW else HANDLER_REGISTRY,
    # )
    # if USE_SIMULATED_WORKFLOW:
    #     logger.warning("⚠️ Workflow system running in SIMULATION mode - no actual DAQ operations")
    #
    # set_workflow_dependencies(workflow_redis_manager, workflow_runner, cfg)

    # ── Initial TELL sync ──
    try:
        daq.sync_current_sample_from_tell(force=True)
    except Exception as e:
        logger.warning(f"Initial sample sync from TELL failed: {e}")

    # ── Wire callbacks ──
    daq.set_face_detection_progress_callback(_push_face_detection_progress)

    logger.info(f"Worker {os.getpid()} initialised successfully.")

    yield  # ── application serves requests here ──

    # Shutdown: add cleanup here if needed
    logger.info(f"Worker {os.getpid()} shutting down.")

app = FastAPI(lifespan=lifespan)
register_exception_handlers(app)

app.include_router(workflow_router)

def _required_recovery_code() -> str:
    """
    Get the required recovery confirmation code from environment variables.

    Returns:
        The recovery confirmation code.

    Raises:
        HTTPException: If AARE_RECOVERY_CODE is not configured.
    """
    code = os.getenv("AARE_RECOVERY_CODE", "").strip()
    if not code:
        logger.error("AARE_RECOVERY_CODE is not configured.")
        raise HTTPException(
            status_code=api_status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Recovery confirmation code is not configured on the server.",
        )
    return code

def _validate_recovery_code(confirmation_code: str) -> None:
    """
    Validate the provided recovery confirmation code against the server configuration.

    Args:
        confirmation_code: The code provided by the user.

    Raises:
        HTTPException: If the code is invalid or not configured.
    """
    expected = _required_recovery_code()
    provided = str(confirmation_code or "").strip()
    if not hmac.compare_digest(provided, expected):
        logger.warning("Invalid recovery confirmation code.")
        raise HTTPException(
            status_code=api_status.HTTP_403_FORBIDDEN,
            detail="Invalid confirmation code.",
        )

def _sample_is_mounted() -> bool:
    """
    Check if a sample is currently mounted on the beamline.

    Returns:
        True if a sample is mounted, False otherwise.
    """
    try:
        return daq.sample is not None
    except Exception:
        return False

def _push_face_detection_progress(payload: dict) -> None:
    """
    Update the global face detection state with new progress information.

    Args:
        payload: Dictionary containing face detection progress data.
    """
    global _face_detection_state
    try:
        next_seq = int(_face_detection_state.get("seq", 0)) + 1
        _face_detection_state = {
            "seq": next_seq,
            **payload,
        }
    except Exception as e:
        logger.warning(f"Failed to update face detection progress: {e}")


async def face_detection_event_stream() -> AsyncGenerator[str, None]:
    """
    Generator for Server-Sent Events (SSE) of face detection progress.

    Yields:
        Formatted SSE data strings.
    """
    last_seq = -1
    try:
        while True:
            state = dict(_face_detection_state)
            seq = int(state.get("seq", 0))
            if seq != last_seq:
                last_seq = seq
                yield f"data: {json.dumps(state, separators=(',', ':'))}\n\n"
            await asyncio.sleep(0.15)
    except asyncio.CancelledError:
        return

@app.post("/token")
async def login(form_data: OAuth2PasswordRequestForm = Depends()):
    """
    Authenticate a user and return an access token.

    Args:
        form_data: OAuth2 password request form containing username and password.

    Returns:
        A dictionary containing the access token and token type.
    """
    data = await run_in_threadpool(auth.authenticate_user, cfg, form_data)
    return {"access_token": data, "token_type": "bearer"}

@app.get("/meta/error-codes")
async def meta_error_codes() -> dict[str, dict[str, str]]:
    """
    Public, stable registry of machine-readable error codes.
    Useful for GUIs, tests, and diagnostics.
    """
    return export_error_codes_grouped()

@app.get("/status")
async def status(token: str = Depends(oauth2_scheme)) -> DAQStatusModel:
    """
    Get the current status of the DAQ system.

    Args:
        token: OAuth2 access token.

    Returns:
        DAQStatusModel containing the current status.
    """
    data = auth.parse_token(token)

    full = daq.status
    active_pgroup = cfg.pgroup
    is_staff = data.staff
    in_allowed_groups = (active_pgroup is not None and active_pgroup in data.pgroups)
    in_ro = is_staff or in_allowed_groups
    sample_pgroup = full.sample.user if full.sample is not None else None
    sample_view_allowed = sample_pgroup in data.pgroups or is_staff

    full.sample             = full.sample if in_ro and sample_view_allowed else None
    full.box                = full.box if in_ro else None
    full.last_best_res      = full.last_best_res if in_ro else None
    full.last_best_b_factor = full.last_best_b_factor if in_ro else None
    full.crystal_size       = full.crystal_size if in_ro else CrystalSize(x=0,y=0,z=0)
    full.session = SessionStatus(
            current_pgroup=cfg.pgroup,
            session=cfg.session_state(data.session),
            staff=data.staff
    )
    return full

@app.get("/beamline/geometry")
async def sample_geometry(
        token: str = Depends(oauth2_scheme),
) -> SampleGeometryModel:
    """
    Get the sample geometry information.

    Args:
        token: OAuth2 access token.

    Returns:
        SampleGeometryModel object.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return daq.status.geom


@app.put("/beamline/omega")
async def omega(val: float, token: str = Depends(oauth2_scheme)):
    """
    Set the omega angle of the goniometer.

    Args:
        val: The target omega angle in degrees.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Setting omega to {val}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.omega = val
    return "OK"

@app.put("/beamline/omega_rel")
async def omega(val: float, token: str = Depends(oauth2_scheme)):
    """
    Move the omega angle relatively.

    Args:
        val: The relative omega movement in degrees.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Moving omega by {val}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.omega_rel(val)
    return "OK"

@app.put("/beamline/front_light")
async def front_light(val: float, token: str = Depends(oauth2_scheme)):
    """
    Set the front light intensity.

    Args:
        val: Light intensity value.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Setting light to {val}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.front_light = val
    return "OK"

@app.put("/beamline/back_light")
async def back_light(val: float, token: str = Depends(oauth2_scheme)):
    """
    Set the back light intensity.

    Args:
        val: Light intensity value.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Setting back light to {val}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.back_light = val
    return "OK"

@app.put("/beamline/zoom")
async def zoom(val: float, token: str = Depends(oauth2_scheme)):
    """
    Set the camera zoom level.

    Args:
        val: Zoom level.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Setting zoom to {val}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.zoom = val
    return "OK"


@app.put("/beamline/smargon")
async def smargon(val: SmargonCoordinate, token: str = Depends(oauth2_scheme)):
    """
    Move the Smargon goniometer to specified coordinates.

    Args:
        val: SmargonCoordinate object.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Setting smargon to {val}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.smargon = val
    return "OK"


@app.post("/beamline/tweak_abr_meas_pos")
async def tweak_abr_meas_pos(val: AerotechCoordinate, token: str = Depends(oauth2_scheme)):
    """
    Tweak the Aerotech measurement position. Staff only.

    Args:
        val: AerotechCoordinate object with tweaks.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.tweak_abr_meas_pos(val)
    return "OK"


@app.post("/beamline/save_abr_meas_pos")
async def save_abr_meas_pos(token: str = Depends(oauth2_scheme)):
    """
    Save the current Aerotech measurement position. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Save abr")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.save_abr_meas_pos()
    return "OK"

@app.post("/beamline/anneal")
async def anneal(time_s: float, token: str = Depends(oauth2_scheme)):
    """
    Perform sample annealing for a specified duration.

    Args:
        time_s: Annealing time in seconds.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Anneal {time_s}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.anneal(time_s)
    return "OK"

@app.post("/beamline/goto_abr_meas_pos")
async def goto_abr_meas_pos(token: str = Depends(oauth2_scheme)):
    """
    Move the Aerotech to the saved measurement position.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Go to ABR meas pos")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.goto_abr_meas_pos()
    return "OK"

@app.post("/beam_mark/add")
async def mark_beam(x: float, y: float, token: str = Depends(oauth2_scheme)):
    """
    Mark the beam position on the camera image. Staff only.

    Args:
        x: X coordinate in pixels.
        y: Y coordinate in pixels.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Marking beam to {x}, {y}")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.mark_beam(x, y)
    return "OK"


@app.post("/beam_mark/clear")
async def clear_beam_mark(token: str = Depends(oauth2_scheme)):
    """
    Clear the beam mark from the camera image. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Clear beam mark")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.clear_mark_beam()
    return "OK"

@app.post("/beamline/beam_center")
async def beam_center(x: float, y: float, token: str = Depends(oauth2_scheme)):
    """
    Set the beam center position. Staff only.

    Args:
        x: X coordinate in pixels.
        y: Y coordinate in pixels.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Beam Center {x}, {y}")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.beam_center = (x, y)
    return "OK"

@app.post("/beamline/beam_size_mm")
async def beam_size_mm(x: float, y: float, token: str = Depends(oauth2_scheme)):
    """
    Set the beam size in millimeters. Staff only.

    Args:
        x: Beam width in mm.
        y: Beam height in mm.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Beam Size {x}, {y}")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.beam_size_mm = Coordinate(x=x, y=y)
    return "OK"

@app.put("/beamline/samcam")
async def samcam_settings(s: SampleCameraSettings, token: str = Depends(oauth2_scheme)):
    """
    Update the sample camera settings (exposure, gain, etc.).

    Args:
        s: SampleCameraSettings object.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"SamCam settings: {s}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.samcam_settings = s
    return "OK"

@app.post("/samcam/autofocus")
async def samcam_autofocus(s: AutofocusSettings, token: str = Depends(oauth2_scheme)):
    """
    Trigger the sample camera autofocus procedure.

    Args:
        s: AutofocusSettings object.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"SamCam AutoFocus")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.auto_focus(s)
    return "OK"

@app.post("/beamline/shutter")
async def shutter(val: bool, token: str = Depends(oauth2_scheme)):
    """
    Open or close the beamline shutter.

    Args:
        val: True to open, False to close.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Shutter {val}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.shutter = val
    return "OK"


async def get_image(token: str = Depends(oauth2_scheme)):
    """
    Get the current camera image as a JPEG stream.

    Args:
        token: OAuth2 access token.

    Returns:
        StreamingResponse containing the JPEG image.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))

    _, encoded_image = cv2.imencode(
        ".jpg", daq.camera_image
    )  # Encodes the image in JPEG format
    image_bytes = io.BytesIO(
        encoded_image.tobytes()
    )  # Convert OpenCV byte format to a file-like object

    return StreamingResponse(image_bytes, media_type="image/jpeg")


# TELL procedures
@app.get("/sample/curr_sample")
async def sample(token: str = Depends(oauth2_scheme)) -> SampleShortInfo:
    """
    Get information about the currently mounted sample.

    Args:
        token: OAuth2 access token.

    Returns:
        SampleShortInfo object.
    """
    token_data = auth.parse_token(token)
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    s = daq.sample
    if token_data.staff or s.user in token_data.pgroups:
        return daq.sample
    else:
        return SampleShortInfo(
            sample_name="Other user sample",
            run_number=0,
            puck_name="",
            dewar_name="",
            db_id=-1,
            pin=s.pin,
            location=s.location
        )

async def park_and_dry(token: str = Depends(oauth2_scheme)):
    """
    Execute the 'park and dry' procedure for the sample changer (TELL).

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with status and message.
    """
    token_data = auth.parse_token(token)
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    daq.park_and_dry()
    return  {
        "ok": True,
        "message": "TELL has been dryed and parked",
    }


@app.post("/sample/mount")
async def mount(dbid: int, token: str = Depends(oauth2_scheme), reference: bool = False):
    """
    Mount a sample from the spreadsheet onto the goniometer.

    Args:
        dbid: Database ID of the sample to mount.
        token: OAuth2 access token.
        reference: Whether to look for the sample in the reference tools list.

    Returns:
        Result of the mounting procedure.
    """
    token_data = auth.parse_token(token)
    logger.debug(f"Mount {dbid}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))

    if reference:
        st = daq.reference_tools
    else:
        st = daq.sample_spreadsheet

    index = -1
    for i in range(len(st.s)):
        if st.s[i].db_id == dbid:
            index = i

    if index == -1:
        raise SampleException(message="Sample not found")

    if not (token_data.staff or st.s[index].user in token_data.pgroups):
        raise UserRightsException(message="Sample belongs to a different user.")

    daq.sample = st.s[index]

    return "OK"

@app.post("/sample/unmount")
async def unmount(token: str = Depends(oauth2_scheme)):
    """
    Unmount the current sample from the goniometer.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Unmount")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.sample = None
    return "OK"


@app.post("/sample/manual")
async def manual(s: SampleShortInfo, token: str = Depends(oauth2_scheme)):
    """
    Manually create or update a sample.

    Args:
        s: SampleShortInfo object.
        token: OAuth2 access token.
    """
    logger.debug(f"Manual Sample {s.sample_name}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    print(f"DB ID prior creating {s.db_id}")
    daq.create_sample(s)
    print(f"DB ID after creating {s.db_id}")

@app.post("/sample/resync")
async def sample_resync(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Manually trigger a resynchronization of the sample information from the changer (TELL).

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with status and message.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.sync_current_sample_from_tell(force=True)
    logger.info("TELL sample cache resynced via API request.")
    return {
        "ok": True,
        "message": "TELL sample cache resynced.",
    }


def get_spreadsheet(data: TokenData) -> SampleShortInfoList:
    """
    Get the sample spreadsheet for the given user/pgroup.

    Args:
        data: TokenData containing user information and pgroups.

    Returns:
        SampleShortInfoList object.
    """
    if data.staff:
        return cfg.spreadsheet
    else:
        return cfg.spreadsheet_pgroup(data.pgroups)

def get_reference_tools() -> SampleShortInfoList:
    """
    Get the list of reference tools.

    Returns:
        SampleShortInfoList object.
    """
    return cfg.reference_tools

async def reference_tools_event_stream() -> AsyncGenerator[str, None]:
    """
    Generator for SSE of reference tools updates.

    Yields:
        JSON string of the reference tools list.
    """
    try:
        while True:
            yield get_reference_tools().model_dump_json()
            await asyncio.sleep(10)
    except asyncio.CancelledError:
        return

async def spreadsheet_event_stream(data: TokenData) -> AsyncGenerator[str, None]:
    """
    Generator for SSE of sample spreadsheet updates.

    Args:
        data: TokenData for filtering the spreadsheet.

    Yields:
        JSON string of the user's sample spreadsheet.
    """
    try:
        while True:
            yield get_spreadsheet(data).model_dump_json()
            await asyncio.sleep(10)
    except asyncio.CancelledError:
        return


@app.get("/sse/spreadsheet")
async def spreadsheet_sse(token: str = Depends(oauth2_scheme)):
    """
    SSE endpoint for sample spreadsheet updates.

    Args:
        token: OAuth2 access token.

    Returns:
        StreamingResponse.
    """
    data = auth.parse_token(token)
    return StreamingResponse(
        spreadsheet_event_stream(data),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Cache-Control"
        }
    )

@app.get("/sse/reference_tools")
async def reference_tools_sse():
    """
    SSE endpoint for reference tools updates.

    Returns:
        StreamingResponse.
    """
    return StreamingResponse(
        reference_tools_event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Cache-Control"
        }
    )

@app.get("/sample/spreadsheet")
async def spreadsheet(token: str = Depends(oauth2_scheme)) -> SampleShortInfoList:
    """
    Get the current sample spreadsheet.

    Args:
        token: OAuth2 access token.

    Returns:
        SampleShortInfoList.
    """
    return get_spreadsheet(auth.parse_token(token))

@app.get("/sample/reference_tools")
async def reference_tools(token: str = Depends(oauth2_scheme)) -> SampleShortInfoList:
    """
    Get the reference tools list.

    Args:
        token: OAuth2 access token.

    Returns:
        SampleShortInfoList.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return get_reference_tools()


# State transitions
@app.post("/state/dewar_exchange")
async def dewar_exchange(token: str = Depends(oauth2_scheme)):
    """
    Transition beamline state to DewarTransfer.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.state = BeamlineStateEnum.DewarTransfer
    return "OK"


@app.post("/state/sample_exchange")
async def dewar_exchange(token: str = Depends(oauth2_scheme)):
    """
    Transition beamline state to SampleExchange.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.state = BeamlineStateEnum.SampleExchange
    return "OK"


@app.post("/state/sample_alignment")
async def sample_alignment(token: str = Depends(oauth2_scheme)):
    """
    Transition beamline state to SampleAlignment.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.state = BeamlineStateEnum.SampleAlignment

@app.post("/state/beam_location")
async def beam_location(token: str = Depends(oauth2_scheme)):
    """
    Transition beamline state to BeamLocation. Staff only.

    Args:
        token: OAuth2 access token.
    """
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.state = BeamlineStateEnum.BeamLocation

@app.post("/state/maintenance")
async def maintenance(token: str = Depends(oauth2_scheme)) -> str:
    """
    Transition beamline state to Maintenance. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    cfg.state = BeamlineStateEnum.Maintenance
    logger.warning("Beamline state set to Maintenance via protected endpoint.")
    return "OK"

@app.post("/access/take_over_beamline")
async def take_over_beamline(payload: RecoveryActionRequest, token: str = Depends(oauth2_scheme)) -> str:
    """
    Take over the beamline session. Staff only.

    Args:
        payload: RecoveryActionRequest with confirmation code.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    _validate_recovery_code(payload.confirmation_code)
    auth.force_current_sesion(cfg, data)
    logger.warning(
        "Beamline session forcefully taken over.",
        extra={"session": getattr(data, "session", None)},
    )
    return "OK"

@app.post("/state/free_beamline")
async def free_beamline(payload: RecoveryActionRequest, token: str = Depends(oauth2_scheme)) -> str:
    """
    Clear the beamline busy flag. Staff only.

    Args:
        payload: RecoveryActionRequest with confirmation code.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    _validate_recovery_code(payload.confirmation_code)
    cfg.state_busy = False
    logger.warning(
        "Beamline busy flag cleared via protected endpoint.",
        extra={"session": getattr(data, "session", None)},
    )
    return "OK"

@app.post("/recovery/recover_beamline")
async def recover_beamline(payload: RecoveryActionRequest, token: str = Depends(oauth2_scheme)) -> dict:
    """
    Recover the beamline from an error state. Staff only.

    Args:
        payload: RecoveryActionRequest with confirmation code.
        token: OAuth2 access token.

    Returns:
        Dictionary with status and previous state information.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    _validate_recovery_code(payload.confirmation_code)

    sample_mounted = _sample_is_mounted()
    prev_state = cfg.state
    prev_busy = cfg.state_busy

    auth.force_current_sesion(cfg, data)
    cfg.state_busy = False
    cfg.state = BeamlineStateEnum.Maintenance

    logger.warning(
        "Beamline recovery action executed.",
        extra={
            "session": getattr(data, "session", None),
            "previous_state": getattr(prev_state, "name", str(prev_state)),
            "previous_busy": prev_busy,
            "sample_mounted": sample_mounted,
        },
    )

    return {
        "ok": True,
        "sample_mounted": sample_mounted,
        "previous_state": getattr(prev_state, "name", str(prev_state)),
        "previous_busy": prev_busy,
        "new_state": BeamlineStateEnum.Maintenance.name,
    }
@app.post("/recovery/unmount_sample")
async def recovery_unmount_sample(payload: RecoveryActionRequest, token: str = Depends(oauth2_scheme)) -> dict:
    """
    Forcefully unmount a sample during recovery. Staff only.

    Args:
        payload: RecoveryActionRequest with confirmation code.
        token: OAuth2 access token.

    Returns:
        Dictionary with status of the recovery unmount action.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    _validate_recovery_code(payload.confirmation_code)

    auth.force_current_sesion(cfg, data)

    if cfg.state_busy:
        raise HTTPException(
            status_code=api_status.HTTP_409_CONFLICT,
            detail="Beamline is busy. Clear or recover the beamline before attempting recovery unmount.",
        )

    status = daq.status
    if not getattr(status, "tell_connected", False):
        raise HTTPException(
            status_code=api_status.HTTP_409_CONFLICT,
            detail=f"TELL is not connected: {getattr(status, 'tell_error', 'unknown error')}",
        )

    sample_mounted = _sample_is_mounted()
    if not sample_mounted:
        return {
            "ok": True,
            "sample_mounted": False,
            "message": "No sample appears to be mounted.",
        }

    prev_state = cfg.state
    daq.recovery_unmount_sample()

    logger.warning(
        "Recovery sample unmount executed.",
        extra={
            "session": getattr(data, "session", None),
            "previous_state": getattr(prev_state, "name", str(prev_state)),
        },
    )

    return {
        "ok": True,
        "sample_mounted": True,
        "previous_state": getattr(prev_state, "name", str(prev_state)),
        "new_state": getattr(cfg.state, "name", str(cfg.state)),
        "message": "Recovery unmount completed.",
    }

# Scans
@app.post("/scan/raster")
async def raster(val: RasterGridRequest, auto: bool = False, token: str = Depends(oauth2_scheme)) -> CompletedRasterGrid:
    """
    Execute a raster scan.

    Args:
        val: RasterGridRequest parameters.
        auto: Whether to perform automated analysis after scan.
        token: OAuth2 access token.

    Returns:
        CompletedRasterGrid object.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    return daq.measure_raster(val, auto)


@app.post("/scan/rotation")
async def rotation(val: RotationScanRequest, token: str = Depends(oauth2_scheme)) -> CompletedRotationScan:
    """
    Execute a rotation scan.

    Args:
        val: RotationScanRequest parameters.
        token: OAuth2 access token.

    Returns:
        CompletedRotationScan object.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    return daq.measure_rotation(val)

@app.post("/scan/auto")
async def auto(s: SampleShortInfo, token: str = Depends(oauth2_scheme)):
    """
    Execute a fully automated measurement sequence for a sample.

    Args:
        s: SampleShortInfo object.
        token: OAuth2 access token.

    Returns:
        Formatted string of the total runtime.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    runtime = daq.measure(s)
    return f"{runtime:0.3f}"


@app.post("/scan/smart_params")
async def set_smart_params(p: SimpleScanParameters, token: str = Depends(oauth2_scheme)) -> str:
    """
    Set the 'smart' scan parameters.

    Args:
        p: SimpleScanParameters object.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    #token_data = auth.parse_token(token)
    #logger.debug(f"{token_data.session} Try to set smart params: {p}")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    cfg.auto_params = p
    return "OK"

@app.post("/scan/cancel")
async def cancel(token: str = Depends(oauth2_scheme)):
    """
    Cancel the currently running scan or automation.

    Args:
        token: OAuth2 access token.
    """
    logger.debug(f"Scan Cancel")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.cancel()

# ALC routines
@app.post("/alc/center_loop")
async def alc_center_loop(token: str = Depends(oauth2_scheme)) -> str:
    """
    Trigger the automated loop centering procedure (ALC).

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"ALC")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.auto_loop_center()
    return "OK"


@app.post("/alc/ml_bounding_box")
async def alc_ml_bounding_box(token: str = Depends(oauth2_scheme)) -> RasterGridRequest | None:
    """
    Request an ML-based bounding box for the sample.

    Args:
        token: OAuth2 access token.

    Returns:
        RasterGridRequest if a box was found, else None.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    return daq.ml_bounding_box()

@app.post("/face_detection/run")
async def face_detection_run(steps: int, step_size: int, token: str = Depends(oauth2_scheme)) -> dict:
    """
    Trigger the face detection sequence.

    Args:
        steps: Number of rotation steps.
        step_size: Degrees per step.
        token: OAuth2 access token.

    Returns:
        Dictionary containing face detection results.
    """
    logger.debug(f"Face detection run: {steps} steps, {step_size} step size")
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    _push_face_detection_progress({
        "running": True,
        "status": "starting",
        "samples": [],
        "height_fit": {},
        "area_fit": {},
    })
    result = daq.face_detection(steps=steps, step_size=step_size)
    return result

@app.get("/sse/face_detection")
async def sse_face_detection(token: str = Depends(oauth2_scheme)):
    """
    SSE endpoint for face detection progress updates.

    Args:
        token: OAuth2 access token.

    Returns:
        StreamingResponse.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return StreamingResponse(
        face_detection_event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Cache-Control"
        }
    )

# Access management
@app.get("/access/pgroup")
async def pgroup(token: str = Depends(oauth2_scheme)) -> str:
    """
    Get the currently active pgroup.

    Args:
        token: OAuth2 access token.

    Returns:
        The pgroup string.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return cfg.pgroup


@app.put("/access/pgroup")
async def set_pgroup(val: str, token: str = Depends(oauth2_scheme)) -> str:
    """
    Set the active pgroup.

    Args:
        val: The pgroup string to set.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    data = auth.parse_token(token)

    holder = cfg.baton_holder
    is_current_holder = holder is not None and holder.session == data.session

    # Staff or current baton holder may change p-group even if it is not currently active.
    # Everyone else must still belong to the active p-group.
    if not (data.staff or is_current_holder):
        auth.check_jwt_ro(cfg, data)

    cfg.pgroup = val
    return "OK"

@app.delete("/access/pgroup")
async def del_pgroup(token: str = Depends(oauth2_scheme)) -> str:
    """
    Clear the active pgroup.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    cfg.pgroup = None
    return "OK"

@app.put("/beamline/commissioning_mode")
async def set_commissioning_mode(val: bool, token: str = Depends(oauth2_scheme)) -> str:
    """
    Set the commissioning mode. Staff only.

    Args:
        val: True to enable commissioning mode, False to disable.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    cfg.commissioning_mode = val
    return "OK"

@app.post("/access/end_session")
async def end_session(token: str = Depends(oauth2_scheme)) -> str:
    """
    End the current session and release the baton.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Try to end Session")
    # End active session will only delete session, if it is equal to token value
    # so no need to check R/W permissions
    token_data = auth.parse_token(token)
    cfg.end_active_session(token_data.session)
    return "OK"


@app.post("/access/force_current_session")
async def force_current_session(token: str = Depends(oauth2_scheme)) -> str:
    """
    Forcefully set the current session as active.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Try to grab session")
    data = auth.parse_token(token)
    # Counterintuitive, this operation requires only R/O permission
    # as this is actually acquiring R/W permissions
    auth.check_jwt_ro(cfg, data)
    auth.force_current_sesion(cfg, data)
    return "OK"

# ========== BATON CONTROL ENDPOINTS ==========

@app.get("/baton/status")
async def baton_status(token: str = Depends(oauth2_scheme)) -> BatonStatus:
    """
    Get the current baton status for the requesting user.

    Args:
        token: OAuth2 access token.

    Returns:
        BatonStatus object.
    """
    data = auth.parse_token(token)
    auth.resolve_baton_timeout_if_needed(cfg)
    return auth.get_baton_status(cfg, data)


@app.post("/baton/request")
async def baton_request(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Request control (baton) of the beamline.

    - If vacant: granted immediately
    - If staff requesting: granted immediately (or queued if busy)
    - If same level: creates pending request with timeout
    - Non-staff cannot request from staff

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with request result.
    """
    logger.debug(cfg.allow_non_staff_request_from_staff)
    data = auth.parse_token(token)
    return auth.request_baton(cfg, data)

@app.post("/baton/respond")
async def baton_respond(accept: bool, token: str = Depends(oauth2_scheme)) -> dict:
    """
    Current baton holder responds to a pending request.

    - accept=true: transfers baton (or queues if busy)
    - accept=false: refuses the request

    Args:
        accept: True to accept, False to refuse.
        token: OAuth2 access token.

    Returns:
        Dictionary with response result.
    """
    data = auth.parse_token(token)
    return auth.respond_to_baton_request(cfg, data, accept)


@app.post("/baton/release")
async def baton_release(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Voluntarily release the baton, making the beamline vacant.

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with release result.
    """
    data = auth.parse_token(token)
    return auth.release_baton(cfg, data)


@app.post("/baton/cancel")
async def baton_cancel(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Cancel your own pending baton request.

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with cancellation result.
    """
    data = auth.parse_token(token)
    return auth.cancel_baton_request(cfg, data)


@app.get("/baton/check_timeout")
async def baton_check_timeout(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Check if a pending request has timed out and process it.
    Called by GUI to poll for timeout completion.

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with timeout check result.
    """
    data = auth.parse_token(token)

    auth.resolve_baton_timeout_if_needed(cfg)

    pending = cfg.pending_baton_request
    if pending is None:
        if cfg.baton_holder and cfg.baton_holder.session == data.session:
            return {"granted": True, "message": "Baton acquired!"}
        queued = cfg.queued_baton_transfer
        if queued and queued.target_session == data.session:
            return {"queued": True, "message": "Transfer queued"}
        return {"no_pending": True}

    if pending.requester_session != data.session:
        return {"not_your_request": True}

    if pending.status == BatonRequestStatus.REFUSED:
        cfg.clear_pending_baton_request()
        return {"refused": True, "message": "Request refused"}

    elapsed = time.time() - pending.created_at
    if elapsed < pending.timeout_seconds:
        return {
            "pending": True,
            "remaining_seconds": pending.timeout_seconds - elapsed
        }

    return auth.request_baton(cfg, data)

@app.put("/access/allow_non_staff_request_from_staff")
async def set_allow_non_staff_request_from_staff(val: bool, token: str = Depends(oauth2_scheme)) -> str:
    """
    Set whether non-staff users can request the baton from staff members. Staff only.

    Args:
        val: True to allow, False to disallow.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    cfg.allow_non_staff_request_from_staff = val
    return "OK"

async def baton_status_event_stream(data: TokenData) -> AsyncGenerator[str, None]:
    """
    SSE stream for baton status updates.

    Args:
        data: TokenData for the requesting user.

    Yields:
        Baton status as SSE data.
    """
    last_status = None
    try:
        while True:
            auth.resolve_baton_timeout_if_needed(cfg)

            new_baton_status = auth.get_baton_status(cfg, data)
            status_json = new_baton_status.model_dump_json()

            if status_json != last_status:
                last_status = status_json
                yield f"data: {status_json}\n\n"

            cfg.process_queued_transfer_if_ready(auth.SESSION_EXPIRE_SECONDS)

            await asyncio.sleep(0.5)
    except asyncio.CancelledError:
        return


@app.get("/sse/baton")
async def sse_baton(token: str = Depends(oauth2_scheme)):
    """
    SSE endpoint for real-time baton status updates.

    Args:
        token: OAuth2 access token.

    Returns:
        StreamingResponse.
    """
    data = auth.parse_token(token)
    return StreamingResponse(
        baton_status_event_stream(data),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Cache-Control"
        }
    )

@app.get("/beamline/settings")
async def get_settings(token: str = Depends(oauth2_scheme)) -> BeamlineSettingsModel:
    """
    Get the current beamline settings. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        BeamlineSettingsModel.
    """
    logger.debug(f"Get settings")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    return cfg.settings


@app.put("/beamline/settings")
async def put_settings(s: BeamlineSettingsModel, token: str = Depends(oauth2_scheme)):
    """
    Update the beamline settings. Staff only.

    Args:
        s: BeamlineSettingsModel object.
        token: OAuth2 access token.
    """
    logger.debug(f"Putting settings: {s}")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    cfg.settings = s

@app.get("/beamline/cryo_settings")
async def get_cryo_settings(token: str = Depends(oauth2_scheme)) -> CryojetSettingsModel:
    """
    Get the current cryojet settings. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        CryojetSettingsModel.
    """
    logger.debug(f"Get Cryo Settings")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    return cfg.cryojet_settings

@app.put("/beamline/cryo_settings")
async def put_cryo_settings(s: CryojetSettingsModel, token: str = Depends(oauth2_scheme)):
    """
    Update the cryojet settings. Staff only.

    Args:
        s: CryojetSettingsModel object.
        token: OAuth2 access token.
    """
    logger.debug(f"Put Cryo Settings: {s}")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    cfg.cryojet_settings = s

@app.put("/access/all_pgroups")
async def get_all_pgroups(token: str = Depends(oauth2_scheme)):
    """
    Get a list of all available pgroups. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        List of pgroup strings.
    """
    logger.debug(f"Get all pgroups")
    auth.check_jwt_staff(cfg, auth.parse_token(token))

    base_path = "/sls/mx/data/"

    now=time.monotonic()
    cached = _all_pgroups_cache.get(base_path)
    if cached is not None:
        items, ts = cached
        if (now - ts) < _ALL_PGROUPS_TTL_S:
            return items

    # Refresh: fast scan using scandir, no regex
    try:
        names: list[str] = []
        with os.scandir(base_path) as it:
            for entry in it:
                name = entry.name
                # quick name filter: p + 5 digits, and directory
                if (
                    len(name) == 6 and
                    name[0] == "p" and
                    name[1:].isdigit() and
                    entry.is_dir(follow_symlinks=False)
                ):
                    names.append(name)
        names.sort(key=lambda d: int(d[1:]))
        _all_pgroups_cache[base_path] = (names, now)
        return names
    except Exception as e:
        logger.error(f"Failed to list pgroups: {e}")
        return []

@app.post("/fluorimeter/spectrum")
async def fluorimeter_spectrum(input: FluorescenceSpectrumParameterModel, token: str = Depends(oauth2_scheme)) -> FluorescenceSpectrumOutputModel:
    """
    Request a fluorescence spectrum measurement.

    Args:
        input: FluorescenceSpectrumParameterModel parameters.
        token: OAuth2 access token.

    Returns:
        FluorescenceSpectrumOutputModel.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    return daq.fluorimeter_take_spectrum(input)

@app.post("/fluorimeter/start")
async def fluorimeter_start(erase: bool = False, token: str = Depends(oauth2_scheme)) -> str:
    """
    Start the fluorimeter measurement.

    Args:
        erase: Whether to erase previous data.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.fluorimeter_start(erase)
    return "OK"

@app.post("/fluorimeter/stop")
async def fluorimeter_stop(token: str = Depends(oauth2_scheme)) -> str:
    """
    Stop the fluorimeter measurement.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.fluorimeter_stop()
    return "OK"

@app.get("/fluorimeter/status")
async def fluorimeter_status(token: str = Depends(oauth2_scheme)) -> int | None:
    """
    Get the current fluorimeter status.

    Args:
        token: OAuth2 access token.

    Returns:
        Fluorimeter status code.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return daq.fluorimeter_status()

@app.get("/fluorimeter/data")
async def fluorimeter_data(token: str = Depends(oauth2_scheme)) -> list[int] | None:
    """
    Get the latest fluorimeter data.

    Args:
        token: OAuth2 access token.

    Returns:
        List of data points.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return daq.fluorimeter_data()

@app.get("/fluorimeter/background")
async def fluorimeter_background(token: str = Depends(oauth2_scheme)) -> list[int] | None:
    """
    Get the fluorimeter background data.

    Args:
        token: OAuth2 access token.

    Returns:
        List of background data points.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return daq.fluorimeter_background()

# Optional: SSE stream for live spectra while acquiring
async def fluorimeter_stream() -> AsyncGenerator[str, None]:
    """
    Generator for SSE of fluorimeter status and data.

    Yields:
        JSON string of fluorimeter status/data as SSE.
    """
    logger.debug("start fluorimeter stream...")
    try:
        while True:
            s = daq.fluorimeter_status()
            d = daq.fluorimeter_data()
            b = daq.fluorimeter_background()
            # Ensure JSON-serializable lists; avoid numpy truthiness
            status = "stopped" if s == 0 else "running"
            logger.debug(f"fluorimeter status: {status}. got data from fluorimeter: {d[0]} and background: {b[0]}")
            data = d.tolist() if hasattr(d, "tolist") else (list(d) if isinstance(d, (tuple, list)) else ([] if d is None else [d]))
            bkg = b.tolist() if hasattr(b, "tolist") else (list(b) if isinstance(b, (tuple, list)) else ([] if b is None else [b]))
            logger.debug(f"fluorimeter data: {data[0]}. fluorimeter background: {bkg[0]}")
            payload = json.dumps({"status": s, "data": data, "background": bkg}, separators=(",", ":"))
            yield f"data: {payload}\n\n"
            if s == 0:
                await asyncio.sleep(0.2)
                logger.debug("fluorimeter stream stopped")
                break
            await asyncio.sleep(0.2)
    except asyncio.CancelledError as e:
        logger.error(f">>> Fluorimeter stream cancelled: {e}")
        return

@app.get("/sse/fluorimeter")
async def sse_fluorimeter(token: str = Depends(oauth2_scheme)):
    """
    SSE endpoint for fluorimeter updates.

    Args:
        token: OAuth2 access token.

    Returns:
        StreamingResponse.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    try:
        daq.fluorimeter_start(erase=False)
    except Exception:
        # Ignore if already running or start not needed
        pass
    return StreamingResponse(
        fluorimeter_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Cache-Control"
        }
    )

@app.post("/samcam/send_screenshot_db")
async def send_screenshot_db(
        filename: str | None = None,
        message: str | None = None,
        token: str = Depends(oauth2_scheme),
) -> str:
    """
    Capture a screenshot and send it to the database.

    Args:
        filename: Optional filename.
        message: Optional message.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    data = auth.parse_token(token)
    auth.check_jwt_rw(cfg, data)
    daq.send_screenshot_db(filename=filename, message=message)
    return "OK"

@app.get("/camera/source")
async def get_camera_source(token: str = Depends(oauth2_scheme)):
    """Get the current camera image source."""
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return {"source": daq._AareDAQ__devs.samcam_source}


@app.post("/camera/source")
async def set_camera_source(use_zmq: bool, token: str = Depends(oauth2_scheme)):
    """Set the camera image source preference."""
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq._AareDAQ__devs.set_camera_source(use_zmq)
    return {"source": daq._AareDAQ__devs.samcam_source}

LOGGING_CONFIG = {
    "version": 1,
    "disable_existing_loggers": False,
    "formatters": {
        "default": {
            "()": "uvicorn.logging.DefaultFormatter",
            "fmt": "%(levelprefix)s %(message)s",
            "use_colors": True,
        },
    },
    "handlers": {
        "default": {
            "formatter": "default",
            "class": "logging.StreamHandler",
            "stream": "ext://sys.stdout",
        },
    },
    "loggers": {
        "uvicorn": {
            "handlers": ["default"],
            "level": "DEBUG",
        },
        "uvicorn.access": {
            "handlers": ["default"],
            "level": "DEBUG",
        },
    },
}

def main():
    # Remove in production!
    urllib3.disable_warnings()

    # Run the application using uvicorn
    uvicorn.run("aare.daq.server:app", host="0.0.0.0", port=5210, workers=2, log_config=LOGGING_CONFIG)


if __name__ == "__main__":
    #start_image_stats_receiver(zmq_url="tcp://129.129.110.12:9089")
    main()
    #stop_image_stats_receiver()