lf_mri_platform/services/reconstructor/api/session.py

from __future__ import annotations

import uuid
from datetime import datetime
from typing import Optional, Dict, List
import os
import uuid
from datetime import datetime
import matplotlib

from .core import STORE_DIR
from .schemas import SessionStatus

matplotlib.use("Agg")


class Session:
    def __init__(self, file_raw: str, file_json: str, file_order: Optional[str], sequence_name: str, digit: str,
                 phase_shift: bool):
        self.session_id = str(uuid.uuid4())
        self.file_raw = file_raw
        self.file_json = file_json
        self.file_order = file_order
        self.sequence_name = sequence_name
        self.digit = digit
        self.phase_shift = phase_shift
        self.status = "queued"
        self.progress = 0.0
        self.message = "queued"
        self.created_at = datetime.utcnow().isoformat()
        self.updated_at = self.created_at
        self.work_dir = os.path.join(STORE_DIR, self.session_id)
        os.makedirs(self.work_dir, exist_ok=True)
        self.result_files: List[str] = []
        self.error_traceback: Optional[str] = None

    def to_status(self) -> SessionStatus:
        return SessionStatus(
            session_id=self.session_id, status=self.status, progress=round(self.progress, 3),
            message=self.message, created_at=self.created_at, updated_at=self.updated_at,
            result_files=[os.path.basename(p) for p in self.result_files], error_traceback=self.error_traceback
        )

    def set(self, *, status: Optional[str] = None, progress: Optional[float] = None, message: Optional[str] = None):
        if status is not None: self.status = status
        if progress is not None: self.progress = max(0.0, min(1.0, progress))
        if message is not None: self.message = message
        self.updated_at = datetime.utcnow().isoformat()


SESSIONS: Dict[str, Session] = {}