"""Runs TotalSegmentator inference on an uploaded volume.

GPU_LOCK is exported (not used internally) so callers can distinguish
"queued, waiting for the GPU" from "actually running" in job status --
see main.py's background task wrapper.
"""

import threading

from totalsegmentator.python_api import totalsegmentator

from . import config

# Single GPU in this deployment: serialize inference across concurrent
# requests to avoid two totalsegmentator() calls fighting over the same
# CUDA context / VRAM.
GPU_LOCK = threading.Lock()


def run_totalsegmentator(input_path, output_path, modality, fast=None):
    """Runs TotalSegmentator and returns output_path.

    modality: a free-text modality string parsed from the user's query
        (e.g. "CT", "CT contrast", "MR", "MR ADC"). Anything starting with
        "CT" routes to task="total"; anything else routes to task="total_mr".
    fast: overrides config.FAST_MODE if given.

    Deliberately does NOT use `roi_subset` to restrict inference to just the
    requested structure, even though that would be faster: roi_subset adds a
    preliminary low-resolution "rough crop" detection pass, and in testing
    that pass silently missed a small structure (spleen) on the total_mr
    task -- roi_subset=["liver"] worked, roi_subset=["spleen"] came back
    completely empty on the exact same volume, while running the full task
    (no roi_subset) found spleen correctly every time. Always segmenting
    everything is slower but has been reliably correct in testing across
    CT/MRI and multiple target organs; the target label is still extracted
    afterward the same way.
    """
    if fast is None:
        fast = config.FAST_MODE

    task = "total" if modality.upper().startswith("CT") else "total_mr"

    totalsegmentator(
        input=str(input_path),
        output=str(output_path),
        task=task,
        fast=fast,
        ml=True,  # single multi-label NIfTI output, not one file per structure
        device="gpu",
        quiet=True,
    )
    return output_path