AareDAQ/src/aare/daq/server.py

import asyncio
import hmac
import importlib
import json
import os
import time
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import Any, ClassVar

import uvicorn
from aarecommon.config.beamline import mx_beamline
from aarecommon.config.logger import get_uvicorn_logging_config, setup_logger
from aarecommon.errors.codes import AareErrorCode, export_error_codes_grouped
from aarecommon.errors.exception_handler import (
    MaintenanceStateException,
    SampleException,
    UserRightsException,
)
from aarecommon.math.coordinate import AerotechCoordinate, Coordinate, SmargonCoordinate
from aarecommon.math.sample_geometry import SampleGeometryModel
from aarecommon.models.auth import BatonRequestStatus, BatonStatus
from aarecommon.models.automation import AutomationProgress
from aarecommon.models.models import (
    AutofocusSettings,
    BeamlineSettingsModel,
    BeamlineStateEnum,
    CryojetSettingsModel,
    CrystalSize,
    DAQStatusModel,
    FluorescenceSpectrumOutputModel,
    FluorescenceSpectrumParameterModel,
    RecoveryActionRequest,
    SampleCameraSettings,
    SampleShortInfo,
    SampleShortInfoList,
    SessionStatus,
    SimpleScanParameters,
    TokenData,
)
from aarecommon.models.raster_grid import CompletedRasterGrid, RasterGridRequest
from aarecommon.models.rotation_scan import CompletedRotationScan, RotationScanRequest
from aareDB import SampleEventType
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi import status as api_status
from fastapi.concurrency import run_in_threadpool
from fastapi.security import OAuth2PasswordBearer
from starlette.responses import StreamingResponse
from uvicorn.workers import UvicornWorker  # deprecated shim, present in pinned 0.34.2

from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch
from aare.beamline_dispatch.protocols import BeamlineDispatch
from aare.daq import auth
from aare.daq.config import BeamlineConfig
from aare.daq.config_model import LocalContactConfigModel
from aare.daq.daq import AareDAQ
from aare.daq.server_exception_handler import register_exception_handlers

logger = setup_logger("aareDAQ")

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

# ── Per-worker state: populated inside the lifespan, after fork ──
cfg: BeamlineConfig
daq: AareDAQ
bl_dispatch: BeamlineDispatch

_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()

_automation_progress_state: dict = {"seq": 0, "progress": None}
_automation_progress_state_lock = asyncio.Lock()


class AareUvicornWorker(UvicornWorker):
    # CONFIG_KWARGS merged last into uvicorn Config (uvicorn/workers.py:69) →
    # keeps our access-log filter + proxy_headers under gunicorn.
    CONFIG_KWARGS: ClassVar[dict[str, Any]] = {
        "log_config": get_uvicorn_logging_config(),
        "proxy_headers": False,
    }


@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 cfg, daq, bl_dispatch

    logger.info(f"Worker {os.getpid()} setting up JWT authentication...")
    bl_dispatch = get_beamline_dispatch()
    auth.init_jwt_key(bl_dispatch.auth)
    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)
    cfg.state = daq.read_current_state_from_bec()

    try:
        cfg.reset_automation_progress()
    except Exception as e:
        logger.warning(f"Failed to reset automation progress Redis keys: {e}", exc_info=True)
    try:
        daq.refresh_detector_metadata_cache()
    except Exception as e:
        logger.warning(f"Initial hardware metadata refresh failed: {e}", exc_info=True)
    # ── 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}", exc_info=True)
    # ── Wire callbacks ──
    daq.set_face_detection_progress_callback(_push_face_detection_progress)
    daq.set_automation_progress_callback(_push_automation_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)


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:
        logger.debug("Could not determine whether a sample is mounted", exc_info=True)
        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}", exc_info=True)


def _get_automation_progress_state() -> dict:
    """
    Read automation progress state from shared config/Redis storage.
    """
    try:
        return cfg.get_automation_progress_state()
    except Exception as e:
        logger.warning(f"Failed to read automation progress from Redis: {e}", exc_info=True)
        return {"seq": 0, "progress": None}


def _push_automation_progress(progress: AutomationProgress) -> None:
    """
    Update the shared automation progress state.

    Args:
        progress: Current automation progress model.
    """
    try:
        state = cfg.set_automation_progress_state(progress)
        logger.info(
            f"[automation_progress push] pid={os.getpid()} "
            f"seq={state.get('seq')} "
            f"current_step={state.get('progress', {}).get('current_step')}"
        )
    except Exception as e:
        logger.warning(
            f"Failed to update automation progress: {type(e).__name__}: {e}", exc_info=True
        )


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


async def automation_progress_event_stream() -> AsyncGenerator[str, None]:
    """
    Generator for Server-Sent Events (SSE) of automation progress.
    """
    last_seq = -1
    try:
        while True:
            state = _get_automation_progress_state()
            seq = int(state.get("seq", 0))
            if seq != last_seq:
                last_seq = seq
                logger.info(
                    f"[automation_progress sse] pid={os.getpid()} seq={seq} has_progress={state.get('progress') is not None}"
                )
                yield f"data: {json.dumps(state, separators=(',', ':'))}\n\n"
            await asyncio.sleep(0.15)
    except asyncio.CancelledError:
        return


@app.post("/token")
async def login(request: Request):
    """
    Authenticate a user and return an access token.

    The request must carry an X-Remote-User header set by the Apache Kerberos
    proxy. The client obtains a token by authenticating via Kerberos (NEGOTIATE)
    against the Apache proxy, which forwards the request with X-Remote-User set.

    Args:
        request: The incoming HTTP request (used to inspect headers and client IP).
        form_data: OAuth2 password request form (unused, required by OAuth2 spec).

    Returns:
        A dictionary containing the access token and token type.
    """
    username = auth.authenticate_from_proxy_header(request)
    data = await run_in_threadpool(auth.authenticate_user, cfg, username)
    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)

    cfg.touch_gui_session(
        session=data.session,
        username=data.sub,
        staff=data.staff,
        expiry_sec=cfg.GUI_SESSION_EXPIRE_SECONDS,
    )

    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 = getattr(full, "sample", None)
    sample_pgroup = sample.user if sample is not None else None
    sample_view_allowed = sample_pgroup in data.pgroups or is_staff

    full.sample = sample if in_ro and sample_view_allowed else None
    full.box = getattr(full, "box", None) if in_ro else None
    full.last_best_res = getattr(full, "last_best_res", None) if in_ro else None
    full.last_best_b_factor = getattr(full, "last_best_b_factor", None) if in_ro else None
    full.crystal_size = (
        getattr(full, "crystal_size", CrystalSize(x=0, y=0, z=0))
        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
    )

    if data.staff:
        full.open_guis = cfg.get_open_gui_sessions()
    else:
        own_gui = cfg.get_gui_session(data.session)
        full.open_guis = [own_gui] if own_gui is not None else []
    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 set_omega_abs(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 set_omega_rel(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.post("/beamline/mono_pitch_scan")
async def mono_pitch_scan(plot: bool = False, token: str = Depends(oauth2_scheme)):
    """
    Run a monochromator pitch scan. Staff only.

    Args:
        plot: Whether to enable plotting in the BEC worker.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Running mono pitch scan (plot={plot})")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.mono_pitch_scan(plot=plot)
    return "OK"


@app.get("/beamline/steer_beam_available")
async def steer_beam_available(token: str = Depends(oauth2_scheme)):
    """
    Check if the beam centring routine is available.

    Args:
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    auth.parse_token(token)
    return daq.steer_beam_available()


@app.post("/beamline/steer_beam")
async def steer_beam(
    x: int | None = None, y: int | None = None, token: str = Depends(oauth2_scheme)
):
    """
    Adjust the mirrors to move the beam to the box. Staff only.

    Args:
        x (int, optional): x-coordinate in sample camera pixels to steer to
        y (int, optional): y-coordinate in sample camera pixels to steer to
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug("Running beam steering routine")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.steer_beam(x, y)
    return "OK"


@app.put("/beamline/change_energy")
async def change_energy(value: float, plot: bool = False, token: str = Depends(oauth2_scheme)):
    """
    Change monochromator energy. Staff only.

    Args:
        value: Target energy.
        plot: Whether to enable plotting in the BEC worker.
        token: OAuth2 access token.

    Returns:
        "OK" on success.
    """
    logger.debug(f"Changing energy to {value} (plot={plot})")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.change_energy(value=value, plot=plot)
    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
    logger.debug(f"smargon set to {daq.smargon}")
    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("Save abr")
    auth.check_jwt_staff(cfg, auth.parse_token(token))
    daq.save_abr_meas_pos()
    return "OK"


@app.post("/beamline/save_beam_location_camera_setting")
async def save_beam_location_camera_setting(token: str = Depends(oauth2_scheme)):
    """
    Persist the camera's current gain/exposure as the beam-location preset for
    the current zoom (stored in Redis, re-applied on future zoom changes).

    Returns:
        "OK" on success.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.save_beam_location_camera_setting()
    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("/smargon/initialize")
async def initialise_smargon(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Initialise Smargon. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with status and message.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    daq.initialise_smargon()
    return {"ok": True, "message": "Smargon initialised."}


@app.post("/bec/load_user_macros")
async def bec_load_user_macros(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Load BEC user macros. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    daq.bec_load_user_macros()
    return {"ok": True, "message": "BEC user macros loaded."}


@app.get("/bec/user_macros")
async def bec_list_all_user_macros(token: str = Depends(oauth2_scheme)) -> list:
    """
    List all BEC user macros. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    return daq.bec_list_all_user_macros()


@app.get("/bec/devices")
async def bec_list_all_devices(token: str = Depends(oauth2_scheme)) -> list:
    """
    List all BEC position devices. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    return daq.bec_list_all_devices()


@app.post("/bec/reinitialise_planner_and_position_devices")
async def bec_reinitialise_planner_and_position_devices(
    method: str = "auto", token: str = Depends(oauth2_scheme)
) -> dict:
    """
    Reinitialise BEC planner and position devices. Staff only.

    Args:
        method:
            "auto" - use beamline default.
            "beamline" - force init_beamline_environment().
            "sample" - force init_se_devices() and planner creation.
        token: OAuth2 access token.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    position_devices = daq.bec_reinitialise_planner_and_position_devices(method=method)
    return {
        "ok": True,
        "method": method,
        "position_devices": position_devices,
        "message": "BEC planner and position devices reinitialised.",
    }


@app.post("/bec/save_current_bs_pos")
async def bec_save_current_bs_pos(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Save the current BEC beamstop work position. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    daq.bec_save_current_bs_pos()
    return {"ok": True, "message": "Saved current BEC beamstop work position."}


@app.post("/bec/save_current_collimator_pos")
async def bec_save_current_collimator_pos(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Save the current BEC collimator work position. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    daq.bec_save_current_collimator_pos()
    return {"ok": True, "message": "Saved current BEC collimator 'in' positions."}


@app.post("/bec/save_current_aerotech_position")
async def bec_save_current_aerotech_position(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Save the current BEC aerotech work position and reload device config. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    daq.bec_save_current_aerotech_position()
    return {"ok": True, "message": "Saved current BEC aerotech work position and reloaded devices."}


def initialise_aerotech(self):
    self._cfg.try_set_busy(timeout=360)
    try:
        self._devs.aerotech.home_aerotech()
        self._cfg.state_busy = False
    except Exception as e:
        self._cfg.state_busy = False
        logger.error(f"Failed to initialise Aerotech: {e}")
        raise


def detector_take_pedestal(self):
    self._cfg.try_set_busy(timeout=360)
    try:
        self._jfjoch.take_pedestal()
        self._cfg.state_busy = False
    except Exception as e:
        self._cfg.state_busy = False
        logger.error(f"Failed to take detector pedestal: {e}")
        raise


def initialise_detector(self):
    self._cfg.try_set_busy(timeout=360)
    try:
        self._jfjoch.initialize()
        self._cfg.state_busy = False
    except Exception as e:
        self._cfg.state_busy = False
        logger.error(f"Failed to initialise detector: {e}")
        raise


@app.get("/about/installed_version")
async def running_version() -> str:
    return importlib.metadata.version("aare")


@app.get("/about/server_file_path")
async def server_file() -> str:
    return __file__


@app.get("/local_contact/simulation_state")
async def local_contact_simulation_state(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Return current runtime simulation state for Local Contact tools. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    return daq.get_runtime_simulation_state()


@app.get("/local_contact/device_state")
async def local_contact_device_state(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Return Local Contact device mode/error state. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    return daq.get_local_contact_device_state()


@app.get("/local_contact/links")
async def local_contact_links(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Return Local Contact web control links. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    return daq.get_local_contact_links()


@app.post("/local_contact/simulate/{device}")
async def local_contact_set_simulation(
    device: str, enabled: bool, token: str = Depends(oauth2_scheme)
) -> dict:
    """
    Enable or disable runtime simulation for a backend device and restart its wrapper. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    result = daq.set_runtime_simulation(device, enabled)
    result["message"] = f"{device} simulation set to {enabled}."
    return result


@app.post("/local_contact/restart/{device}")
async def local_contact_restart_device(device: str, token: str = Depends(oauth2_scheme)) -> dict:
    """
    Restart a Local Contact backend wrapper. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)

    device = str(device).strip().lower()
    if device == "bec":
        result = daq.restart_bec_worker()
    elif device == "detector":
        result = daq.restart_detector()
    elif device == "tell":
        result = daq.restart_tell()
    elif device == "aerotech":
        result = daq.restart_aerotech()
    elif device == "smargon":
        result = daq.restart_smargon()
    else:
        raise HTTPException(
            status_code=api_status.HTTP_400_BAD_REQUEST,
            detail="Unknown restart device. Expected one of: bec, detector, tell, aerotech, smargon.",
        )

    result["message"] = f"{device} backend restarted."
    return result


@app.post("/local_contact/resync/detector_metadata")
async def local_contact_resync_detector_metadata(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Refresh cached detector metadata and DTZ limits. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)

    payload = daq.refresh_detector_metadata_cache()
    return {"ok": True, "message": "Hardware metadata cache resynced.", "payload": payload}


@app.get("/local_contact/config")
async def local_contact_config(token: str = Depends(oauth2_scheme)) -> LocalContactConfigModel:
    """
    Return Local Contact config values. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    return daq.get_local_contact_config()


@app.put("/local_contact/config")
async def local_contact_set_config(
    payload: LocalContactConfigModel, token: str = Depends(oauth2_scheme)
) -> LocalContactConfigModel:
    """
    Update Local Contact config values. Staff only.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    return daq.set_local_contact_config(payload)


@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("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("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.put("/beamline/autoexposure")
async def samcam_autoexposure(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.
    """
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    daq.auto_exposure()
    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("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"


# 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,
        )


@app.post("/tell/park_and_dry")
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.
    """
    # TODO: add unmount toggle, and combine aprk_and_dry and tell_dry to one command that takes park and unmount as signals
    auth.check_jwt_rw(cfg, auth.parse_token(token))
    logger.debug("Executing unmount, dry and park")
    daq.park_and_dry(park=True, unmount=True)
    return {"ok": True, "message": "TELL has been dried and parked"}


@app.post("/tell/dry")
async def tell_dry(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Execute a TELL dry cycle. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with status and message.
    """
    data = auth.parse_token(token)
    auth.check_jwt_rw(data)
    logger.debug("Executing dry and return to dewar")
    daq.park_and_dry(park=False, unmount=False)
    return {"ok": True, "message": "TELL dry cycle completed."}


@app.post("/tell/toggle_blower")
async def tell_toggle_blower(token: str = Depends(oauth2_scheme)) -> dict:
    """
    Toggle the TELL blower. Staff only.

    Args:
        token: OAuth2 access token.

    Returns:
        Dictionary with status and message.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    daq.blower_control()
    return {"ok": True, "message": "TELL blower toggled."}


@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))

    daq.check_tell_mount_start_conditions()

    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("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))
    logger.debug(f"DB ID prior creating {s.db_id}")
    daq.create_sample(s)
    logger.debug(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 sample_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/beamstop_alignment")
async def beamstop_alignment(token: str = Depends(oauth2_scheme)):
    """
    Transition beamline state to BeamstopAlignment.

    Args:
        token: OAuth2 access token.

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


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

    Args:
        token: OAuth2 access token.

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


@app.post("/state/data_collection")
async def data_collection(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.DataCollection
    return "OK"


@app.post("/state/robot_sample_exchange")
async def robot_sample_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.RobotSampleExchange
    return "OK"


@app.post("/state/xray_fluorescence")
async def xray_fluorescence(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.XrayFluorescence


@app.post("/state/xtal_snapshot")
async def xtal_snapshot(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.XtalSnapshot


@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/force_clear_busy")
async def force_clear_busy(
    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/force_maintenance_state")
async def force_maintenance_state(
    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 - forced maintenance state.",
        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_center: 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_center)


@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))
    daq.check_tell_mount_start_conditions()
    try:
        runtime = daq.measure(s)
    except Exception as e:
        logger.exception("Critical automation failure in /scan/auto")
        raise HTTPException(
            status_code=api_status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail={
                "code": AareErrorCode.INTERNAL_ERROR.value,
                "message": str(e) or "Critical automation failure",
            },
        ) from e
    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.
    """
    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("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("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",
        },
    )


@app.get("/sse/automation_progress")
async def sse_automation_progress(token: str = Depends(oauth2_scheme)):
    """
    SSE endpoint for DAQ automation progress updates.

    Args:
        token: OAuth2 access token.

    Returns:
        StreamingResponse.
    """
    auth.check_jwt_ro(cfg, auth.parse_token(token))
    return StreamingResponse(
        automation_progress_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.get("/admin/gui_sessions")
async def get_gui_sessions(token: str = Depends(oauth2_scheme)) -> list[dict]:
    """
    Staff-only list of currently active GUIs.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)
    return [item.model_dump() for item in cfg.get_open_gui_sessions()]


@app.post("/admin/gui_sessions/{session_id}/request_close")
async def request_gui_close(
    session_id: int, grace_seconds: int = 60, token: str = Depends(oauth2_scheme)
) -> dict:
    """
    Staff-only request for a remote GUI to close gracefully.
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)

    payload = cfg.request_gui_close(
        session=session_id, requested_by=data.sub, grace_seconds=grace_seconds
    )
    if payload is None:
        raise HTTPException(
            status_code=api_status.HTTP_404_NOT_FOUND, detail="GUI session not found."
        )

    return {"ok": True, "session": session_id, "grace_seconds": grace_seconds}


@app.delete("/admin/gui_sessions/{session_id}")
async def force_remove_gui_session(session_id: int, token: str = Depends(oauth2_scheme)) -> dict:
    """
    Staff-only hard removal of a GUI session from Redis.

    This is intended for stuck/crashed GUIs and clears:
    - GUI presence entry
    - active session/baton if the same session holds it
    """
    data = auth.parse_token(token)
    auth.check_jwt_staff_only(data)

    cfg.end_active_session(session_id)
    cfg.remove_gui_session(session_id)

    return {"ok": True, "session": session_id, "message": "GUI session force removed."}


@app.post("/admin/gui_sessions/{session_id}/interaction")
async def update_gui_interaction(session_id: int, token: str = Depends(oauth2_scheme)) -> dict:
    """
    GUI-side activity heartbeat.
    This is distinct from /status so idle timeout ignores status polling.
    """
    data = auth.parse_token(token)
    if data.session != session_id:
        raise HTTPException(
            status_code=api_status.HTTP_403_FORBIDDEN,
            detail="Cannot update interaction for another session.",
        )

    payload = cfg.update_gui_interaction(session=session_id, last_interaction_ts=time.time())
    if payload is None:
        raise HTTPException(
            status_code=api_status.HTTP_404_NOT_FOUND, detail="GUI session not found."
        )

    return {"ok": True}


@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("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)
    cfg.remove_gui_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("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("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("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("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:
        logger.exception("Failed to list pgroups")
        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
        logger.debug("Fluorimeter already running or start not needed", exc_info=True)
    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"


async def send_message_db(
    db_id: int,
    event_type: SampleEventType,
    comment: str | None = None,
    token: str = Depends(oauth2_scheme),
):
    data = auth.parse_token(token)
    auth.check_jwt_rw(cfg, data)
    daq.send_screenshot_db(db_id=db_id, event_type=event_type, comment=comment)
    return "OK"


@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.")
    raise MaintenanceStateException("Beamline was set to Maintenance. Automation must stop.")


def main():
    uvicorn.run(
        "aare.daq.server:app",
        host="127.0.0.1",
        port=5210,
        workers=4,
        proxy_headers=False,
        log_config=get_uvicorn_logging_config(),
        timeout_worker_healthcheck=30,
    )


if __name__ == "__main__":
    main()