session_control.py

from datetime import datetime
from logging import getLogger
from pathlib import Path
from typing import Dict, List, Optional

from fastapi import APIRouter, Depends
from fastapi.responses import FileResponse
from ispyb.sqlalchemy import AutoProcProgram as ISPyBAutoProcProgram
from pydantic import BaseModel
from sqlalchemy import func
from sqlmodel import select

try:
    from smartem_backend.api_client import SmartEMAPIClient
    from smartem_common.schemas import AtlasData

    SMARTEM_ACTIVE = True
except ImportError:
    SMARTEM_ACTIVE = False

import murfey.server.prometheus as prom
from murfey.server import _transport_object
from murfey.server.api.auth import (
    MurfeySessionIDInstrument as MurfeySessionID,
    validate_instrument_token,
)
from murfey.server.api.session_shared import (
    find_upstream_visits as _find_upstream_visits,
    gather_upstream_files as _gather_upstream_files,
    gather_upstream_tiffs as _gather_upstream_tiffs,
    get_foil_hole as _get_foil_hole,
    get_foil_holes_from_grid_square as _get_foil_holes_from_grid_square,
    get_grid_squares as _get_grid_squares,
    get_grid_squares_from_dcg as _get_grid_squares_from_dcg,
    get_tiff_file as _get_tiff_file,
    get_upstream_file as _get_upstream_file,
    remove_session_by_id,
)
from murfey.server.ispyb import DB as ispyb_db, get_all_ongoing_visits
from murfey.server.murfey_db import murfey_db
from murfey.util import sanitise
from murfey.util.config import get_machine_config
from murfey.util.db import (
    AutoProcProgram,
    ClientEnvironment,
    DataCollection,
    DataCollectionGroup,
    FoilHole,
    GridSquare,
    Movie,
    ProcessingJob,
    RsyncInstance,
    Session,
)
from murfey.util.models import (
    BatchPositionParameters,
    ClientInfo,
    FoilHoleParameters,
    GridSquareParameters,
    RsyncerInfo,
    SearchMapParameters,
    UpstreamFileRequestInfo,
    Visit,
)
from murfey.workflows.spa.atlas import atlas_jpg_from_mrc
from murfey.workflows.spa.flush_spa_preprocess import (
    register_foil_hole as _register_foil_hole,
    register_grid_square as _register_grid_square,
)
from murfey.workflows.tomo.tomo_metadata import (
    register_batch_position_in_database,
    register_search_map_in_database,
)

logger = getLogger("murfey.server.api.session_control")

router = APIRouter(
    prefix="/session_control",
    dependencies=[Depends(validate_instrument_token)],
    tags=["Session Control: General"],
)


@router.get("/time")
async def get_current_timestamp():
    return {"timestamp": datetime.now().timestamp()}


@router.get("/instruments/{instrument_name}/machine")
def machine_info_by_instrument(instrument_name: str):
    return get_machine_config(instrument_name)[instrument_name]


@router.get("/new_client_id/")
async def new_client_id(db=murfey_db):
    clients = db.exec(select(ClientEnvironment)).all()
    if not clients:
        return {"new_id": 0}
    sorted_ids = sorted([c.client_id for c in clients])
    return {"new_id": sorted_ids[-1] + 1}


@router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit])
def get_current_visits(instrument_name: str, db=ispyb_db):
    logger.debug(
        f"Received request to look up ongoing visits for {sanitise(instrument_name)}"
    )
    return get_all_ongoing_visits(instrument_name, db)


class SessionInfo(BaseModel):
    session_id: Optional[int] = None
    session_name: str = ""
    rescale: bool = True


@router.post("/instruments/{instrument_name}/clients/{client_id}/session")
def link_client_to_session(
    instrument_name: str, client_id: int, sess: SessionInfo, db=murfey_db
):
    sid = sess.session_id
    if sid is None:
        s = Session(name=sess.session_name, instrument_name=instrument_name)
        db.add(s)
        db.commit()
        sid = s.id
    client = db.exec(
        select(ClientEnvironment).where(ClientEnvironment.client_id == client_id)
    ).one()
    client.session_id = sid
    db.add(client)
    db.commit()
    db.close()
    return sid


@router.post("/visits/{visit_name}")
def register_client_to_visit(visit_name: str, client_info: ClientInfo, db=murfey_db):
    client_env = db.exec(
        select(ClientEnvironment).where(ClientEnvironment.client_id == client_info.id)
    ).one()
    session = db.exec(select(Session).where(Session.id == client_env.session_id)).one()
    if client_env:
        client_env.visit = visit_name
        db.add(client_env)
        db.commit()
    if session:
        session.visit = visit_name
        db.add(session)
        db.commit()
    db.close()
    return client_info


@router.get("/sessions")
async def get_sessions(db=murfey_db):
    sessions = db.exec(select(Session)).all()
    clients = db.exec(select(ClientEnvironment)).all()
    res = []
    for sess in sessions:
        r = {"session": sess, "clients": []}
        for cl in clients:
            if cl.session_id == sess.id:
                r["clients"].append(cl)
        res.append(r)
    return res


@router.delete("/sessions/{session_id}")
def remove_session(session_id: MurfeySessionID, db=murfey_db):
    remove_session_by_id(session_id, db)


@router.get("/sessions/{session_id}/data_collection_groups")
def get_dc_groups(
    session_id: MurfeySessionID, db=murfey_db
) -> Dict[str, DataCollectionGroup]:
    data_collection_groups = db.exec(
        select(DataCollectionGroup).where(DataCollectionGroup.session_id == session_id)
    ).all()
    return {dcg.tag: dcg for dcg in data_collection_groups}


@router.post("/sessions/{session_id}/successful_processing")
def register_processing_success_in_ispyb(
    session_id: MurfeySessionID, db=ispyb_db, murfey_db=murfey_db
):
    collected_ids = murfey_db.exec(
        select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram)
        .where(DataCollectionGroup.session_id == session_id)
        .where(DataCollection.dcg_id == DataCollectionGroup.id)
        .where(ProcessingJob.dc_id == DataCollection.id)
        .where(AutoProcProgram.pj_id == ProcessingJob.id)
    ).all()
    appids = [c[3].id for c in collected_ids]
    if _transport_object:
        if db is not None:
            apps = db.query(ISPyBAutoProcProgram).filter(
                ISPyBAutoProcProgram.autoProcProgramId.in_(appids)
            )
            for updated in apps:
                updated.processingStatus = True
                _transport_object.do_update_processing_status(updated)


@router.get("/num_movies")
def count_number_of_movies(db=murfey_db) -> Dict[str, int]:
    res = db.exec(
        select(Movie.tag, func.count(Movie.murfey_id)).group_by(Movie.tag)
    ).all()
    return {r[0]: r[1] for r in res}


class PostInfo(BaseModel):
    router_name: str
    function_name: str
    data: dict
    kwargs: dict


@router.post("/instruments/{instrument_name}/failed_client_post")
def failed_client_post(instrument_name: str, post_info: PostInfo):
    zocalo_message = {
        "register": "failed_client_post",
        "router_name": post_info.router_name,
        "function_name": post_info.function_name,
        "data": post_info.data,
        "kwargs": post_info.kwargs,
    }
    if _transport_object:
        _transport_object.send(_transport_object.feedback_queue, zocalo_message)


@router.post("/sessions/{session_id}/rsyncer")
def register_rsyncer(session_id: int, rsyncer_info: RsyncerInfo, db=murfey_db):
    visit_name = db.exec(select(Session).where(Session.id == session_id)).one().visit
    rsync_instance = RsyncInstance(
        source=rsyncer_info.source,
        session_id=rsyncer_info.session_id,
        transferring=rsyncer_info.transferring,
        destination=rsyncer_info.destination,
        tag=rsyncer_info.tag,
    )
    db.add(rsync_instance)
    db.commit()
    db.close()
    prom.seen_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).set(0)
    prom.transferred_files.labels(
        rsync_source=rsyncer_info.source, visit=visit_name
    ).set(0)
    prom.transferred_files_bytes.labels(
        rsync_source=rsyncer_info.source, visit=visit_name
    ).set(0)
    prom.seen_data_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).set(
        0
    )
    prom.transferred_data_files.labels(
        rsync_source=rsyncer_info.source, visit=visit_name
    ).set(0)
    prom.transferred_data_files_bytes.labels(
        rsync_source=rsyncer_info.source, visit=visit_name
    ).set(0)
    prom.skipped_files.labels(rsync_source=rsyncer_info.source, visit=visit_name).set(0)
    return rsyncer_info


@router.get("/sessions/{session_id}/rsyncers", response_model=List[RsyncInstance])
def get_rsyncers_for_session(session_id: MurfeySessionID, db=murfey_db):
    rsync_instances = db.exec(
        select(RsyncInstance).where(RsyncInstance.session_id == session_id)
    )
    return rsync_instances.all()


class StringOfPathModel(BaseModel):
    path: str


@router.post("/sessions/{session_id}/rsyncer_stopped")
def register_stopped_rsyncer(
    session_id: int, rsyncer_source: StringOfPathModel, db=murfey_db
):
    rsyncer = db.exec(
        select(RsyncInstance)
        .where(RsyncInstance.session_id == session_id)
        .where(RsyncInstance.source == rsyncer_source.path)
    ).one()
    rsyncer.transferring = False
    db.add(rsyncer)
    db.commit()


@router.post("/sessions/{session_id}/rsyncer_started")
def register_restarted_rsyncer(
    session_id: int, rsyncer_source: StringOfPathModel, db=murfey_db
):
    rsyncer = db.exec(
        select(RsyncInstance)
        .where(RsyncInstance.session_id == session_id)
        .where(RsyncInstance.source == rsyncer_source.path)
    ).one()
    rsyncer.transferring = True
    db.add(rsyncer)
    db.commit()


@router.delete("/sessions/{session_id}/rsyncer")
def delete_rsyncer(session_id: int, rsyncer_source: StringOfPathModel, db=murfey_db):
    try:
        rsync_instance = db.exec(
            select(RsyncInstance)
            .where(RsyncInstance.session_id == session_id)
            .where(RsyncInstance.source == str(rsyncer_source.path))
        ).one()
        db.delete(rsync_instance)
        db.commit()
    except Exception:
        logger.error(
            f"Failed to delete rsyncer for source directory {sanitise(str(rsyncer_source.path))!r} "
            f"in session {session_id}.",
            exc_info=True,
        )


spa_router = APIRouter(
    prefix="/session_control/spa",
    dependencies=[Depends(validate_instrument_token)],
    tags=["Session Control: SPA"],
)


@spa_router.get("/sessions/{session_id}/grid_squares")
def get_grid_squares(session_id: MurfeySessionID, db=murfey_db):
    return _get_grid_squares(session_id, db)


@spa_router.get("/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares")
def get_grid_squares_from_dcg(
    session_id: MurfeySessionID, dcgid: int, db=murfey_db
) -> List[GridSquare]:
    return _get_grid_squares_from_dcg(session_id, dcgid, db)


@spa_router.get(
    "/sessions/{session_id}/data_collection_groups/{dcgid}/grid_squares/{gsid}/foil_holes"
)
def get_foil_holes_from_grid_square(
    session_id: MurfeySessionID, dcgid: int, gsid: int, db=murfey_db
) -> List[FoilHole]:
    return _get_foil_holes_from_grid_square(session_id, dcgid, gsid, db)


@spa_router.get("/sessions/{session_id}/foil_hole/{fh_name}")
def get_foil_hole(
    session_id: MurfeySessionID, fh_name: int, db=murfey_db
) -> Dict[str, int]:
    return _get_foil_hole(session_id, fh_name, db)


class AtlasRegistration(BaseModel):
    name: str
    acqusition_uuid: str


@spa_router.post("/sessions/{session_id}/register_atlas")
def register_atlas(
    session_id: MurfeySessionID,
    atlas_registration_data: AtlasRegistration,
    db=murfey_db,
):
    if SMARTEM_ACTIVE:
        session = db.exec(select(Session).where(Session.id == session_id)).one()
        machine_config = get_machine_config(session.instrument_name)[
            session.instrument_name
        ]
        if machine_config.smartem_api_url:
            smartem_client = SmartEMAPIClient(
                base_url=machine_config.smartem_api_url, logger=logger
            )
            possible_grids = smartem_client.get_acquisition_grids(
                atlas_registration_data.acqusition_uuid
            )
            grid_uuid = None
            for grid in possible_grids:
                if grid.name == atlas_registration_data.name.replace("_atlas", ""):
                    grid_uuid = grid.uuid
                    break
            if grid_uuid is not None:
                atlas_data = AtlasData(
                    id=atlas_registration_data.name,
                    acquisition_data=datetime.now(),
                    storage_folder="",
                    name=atlas_registration_data.name,
                    tiles=[],
                    gridsquare_positions=None,
                    grid_uuid=grid_uuid,
                )
                smartem_client.create_grid_atlas(atlas_data)
    else:
        logger.info("smartem deactivated so did not register atlas")


@spa_router.post("/sessions/{session_id}/make_atlas_jpg")
def make_atlas_jpg(
    session_id: MurfeySessionID, atlas_mrc: StringOfPathModel, db=murfey_db
):
    logger.debug(
        f"Received request to create JPG image of atlas {sanitise(atlas_mrc.path)!r}"
    )
    session = db.exec(select(Session).where(Session.id == session_id)).one()
    return atlas_jpg_from_mrc(
        session.instrument_name, session.visit, Path(atlas_mrc.path)
    )


@spa_router.post("/sessions/{session_id}/grid_square/{gsid}")
def register_grid_square(
    session_id: MurfeySessionID,
    gsid: int,
    grid_square_params: GridSquareParameters,
    db=murfey_db,
):
    return _register_grid_square(session_id, gsid, grid_square_params, db)


@spa_router.post("/sessions/{session_id}/grid_square/{gs_name}/foil_hole")
def register_foil_hole(
    session_id: MurfeySessionID,
    gs_name: int,
    foil_hole_params: FoilHoleParameters,
    db=murfey_db,
):
    logger.info(
        f"Registering foil hole {foil_hole_params.name} with position {(foil_hole_params.x_location, foil_hole_params.y_location)}"
    )
    return _register_foil_hole(session_id, gs_name, foil_hole_params, db)


tomo_router = APIRouter(
    prefix="/session_control/tomo",
    dependencies=[Depends(validate_instrument_token)],
    tags=["Session Control: CryoET"],
)


@tomo_router.post("/sessions/{session_id}/search_map/{sm_name}")
def register_search_map(
    session_id: MurfeySessionID,
    sm_name: str,
    search_map_params: SearchMapParameters,
    db=murfey_db,
):
    return register_search_map_in_database(session_id, sm_name, search_map_params, db)


@tomo_router.post("/sessions/{session_id}/batch_position/{batch_name}")
def register_batch_position(
    session_id: MurfeySessionID,
    batch_name: str,
    batch_params: BatchPositionParameters,
    db=murfey_db,
):
    return register_batch_position_in_database(session_id, batch_name, batch_params, db)


correlative_router = APIRouter(
    prefix="/session_control/correlative",
    dependencies=[Depends(validate_instrument_token)],
    tags=["Session Control: Correlative Imaging"],
)


@correlative_router.get("/sessions/{session_id}/upstream_visits")
async def find_upstream_visits(session_id: MurfeySessionID, db=murfey_db):
    return _find_upstream_visits(session_id=session_id, db=db)


@correlative_router.get(
    "/visits/{visit_name}/sessions/{session_id}/upstream_file_paths"
)
async def gather_upstream_files(
    visit_name: str,
    session_id: MurfeySessionID,
    upstream_file_request: UpstreamFileRequestInfo,
    db=murfey_db,
):
    return _gather_upstream_files(
        session_id=session_id,
        upstream_instrument=upstream_file_request.upstream_instrument,
        upstream_visit_path=upstream_file_request.upstream_visit_path,
        db=db,
    )


@correlative_router.get(
    "/visits/{visit_name}/sessions/{session_id}/upstream_file/{upstream_file_path:path}"
)
async def get_upstream_file(
    visit_name: str,
    session_id: MurfeySessionID,
    upstream_file_path: str,
    db=murfey_db,
):
    upstream_file = _get_upstream_file(upstream_file_path)
    return (
        FileResponse(path=upstream_file) if upstream_file is not None else upstream_file
    )


@correlative_router.get(
    "/visits/{visit_name}/sessions/{session_id}/upstream_tiff_paths"
)
async def gather_upstream_tiffs(visit_name: str, session_id: int, db=murfey_db):
    return _gather_upstream_tiffs(visit_name=visit_name, session_id=session_id, db=db)


@correlative_router.get(
    "/visits/{visit_name}/sessions/{session_id}/upstream_tiff/{tiff_path:path}"
)
async def get_tiff_file(visit_name: str, session_id: int, tiff_path: str, db=murfey_db):
    tiff_file = _get_tiff_file(
        visit_name=visit_name, session_id=session_id, tiff_path=tiff_path, db=db
    )
    return FileResponse(path=tiff_file) if tiff_file is not None else tiff_file