context.py

from __future__ import annotations

import logging
from importlib.metadata import entry_points
from pathlib import Path
from typing import Any, NamedTuple, OrderedDict

import xmltodict

from murfey.client.instance_environment import MurfeyInstanceEnvironment, SampleInfo
from murfey.util.client import capture_post

logger = logging.getLogger("murfey.client.context")


def _file_transferred_to(
    environment: MurfeyInstanceEnvironment,
    source: Path,
    file_path: Path,
    rsync_basepath: Path,
):
    if environment.visit in environment.default_destinations[source]:
        return (
            rsync_basepath
            / Path(environment.default_destinations[source])
            / file_path.relative_to(source)  # need to strip out the rsync_module name
        )
    return (
        rsync_basepath
        / Path(environment.default_destinations[source])
        / environment.visit
        / file_path.relative_to(source)
    )


def _get_source(file_path: Path, environment: MurfeyInstanceEnvironment) -> Path | None:
    possible_sources = []
    for s in environment.sources:
        if file_path.is_relative_to(s):
            possible_sources.append(s)
    if not possible_sources:
        return None
    elif len(possible_sources) == 1:
        return possible_sources[0]
    source = possible_sources[0]
    for extra_source in possible_sources[1:]:
        if extra_source.is_relative_to(source):
            source = extra_source
    return source


def _atlas_destination(
    environment: MurfeyInstanceEnvironment,
    source: Path,
    rsync_basepath: Path,
) -> Path:
    for i, destination_part in enumerate(
        Path(environment.default_destinations[source]).parts
    ):
        if destination_part == environment.visit:
            return rsync_basepath / "/".join(
                Path(environment.default_destinations[source]).parent.parts[: i + 1]
            )
    return (
        rsync_basepath
        / Path(environment.default_destinations[source]).parent
        / environment.visit
    )


def ensure_dcg_exists(
    collection_type: str,
    metadata_source: Path,
    environment: MurfeyInstanceEnvironment,
    machine_config: dict,
    token: str,
) -> str | None:
    """Create  a data collection group"""
    if collection_type == "tomo":
        experiment_type_id = 36
        session_file = metadata_source / "Session.dm"
        source_visit_dir = metadata_source.parent
    elif collection_type == "spa":
        experiment_type_id = 37
        # For SPA the metadata source sent should include the Images-Disc
        session_file = metadata_source.parent / "EpuSession.dm"
        source_visit_dir = metadata_source.parent.parent
        for h in entry_points(group="murfey.hooks"):
            try:
                if h.name == "get_epu_session_metadata":
                    h.load()(
                        destination_dir=session_file.parent,
                        environment=environment,
                        token=token,
                    )
            except Exception as e:
                logger.warning(f"Get EPU session hook failed: {e}")
    elif collection_type == "sxt":
        experiment_type_id = 47
        session_file = metadata_source / "Session.dm"
        source_visit_dir = metadata_source.parent
    else:
        logger.error(f"Unknown collection type {collection_type}")
        return None

    if not session_file.is_file():
        logger.warning(f"Cannot find session file {str(session_file)}")
        dcg_tag = "/".join(
            [part for part in metadata_source.parts if part != environment.visit]
        ).replace("//", "/")
        dcg_data = {
            "experiment_type_id": experiment_type_id,
            "tag": dcg_tag,
        }
    else:
        with open(session_file, "r") as session_xml:
            session_data = xmltodict.parse(session_xml.read())

        if collection_type == "tomo":
            windows_path = session_data["TomographySession"]["AtlasId"]
        else:
            windows_path = session_data["EpuSessionXml"]["Samples"]["_items"][
                "SampleXml"
            ][0]["AtlasId"]["#text"]

        logger.info(f"Windows path to atlas metadata found: {windows_path}")
        if not windows_path:
            logger.warning("No atlas metadata path found")
            return None
        if environment.visit in windows_path.split("\\"):
            # Case of atlas in the correct location
            visit_index = windows_path.split("\\").index(environment.visit)
            partial_path = "/".join(windows_path.split("\\")[visit_index + 1 :])
            logger.info(
                f"Partial Linux path successfully constructed from Windows path: {partial_path}"
            )
        else:
            # Atlas not in visit, so come up with where it should have been
            # Assumes /structure/to/Supervisor/Sample/Atlas/Atlas.dm
            partial_path = "atlas/" + "/".join(windows_path.split("\\")[-4:])
            logger.info(f"Partial Linux path estimated: {partial_path}")

        logger.info(
            f"Looking for atlas XML file in metadata directory {str((source_visit_dir / partial_path).parent)}"
        )

        dcg_tag = "/".join(
            [part for part in metadata_source.parts if part != environment.visit]
        ).replace("//", "/")

        for p in partial_path.split("/"):
            if p.startswith("Sample"):
                sample = int(p.replace("Sample", ""))
                break
        else:
            logger.warning(f"Sample could not be identified for {metadata_source}")
            return None
        environment.samples[metadata_source] = SampleInfo(
            atlas=Path(partial_path), sample=sample
        )

        if atlas_xml_search := list(
            (source_visit_dir / partial_path).parent.glob("Atlas_*.xml")
        ):
            atlas_xml_path = atlas_xml_search[0]
            logger.info(f"Atlas XML path {str(atlas_xml_path)} found")
            with open(atlas_xml_path, "rb") as atlas_xml:
                atlas_xml_data = xmltodict.parse(atlas_xml)
                atlas_original_pixel_size = float(
                    atlas_xml_data["MicroscopeImage"]["SpatialScale"]["pixelSize"]["x"][
                        "numericValue"
                    ]
                )
            # need to calculate the pixel size of the downscaled image
            atlas_pixel_size = atlas_original_pixel_size * 7.8
            logger.info(f"Atlas image pixel size determined to be {atlas_pixel_size}")

            dcg_data = {
                "experiment_type_id": experiment_type_id,
                "tag": dcg_tag,
                "atlas": str(
                    _atlas_destination(
                        environment,
                        session_file.parent,
                        Path(machine_config.get("rsync_basepath", "")),
                    )
                    / environment.samples[metadata_source].atlas.parent
                    / atlas_xml_path.with_suffix(".jpg").name
                ).replace("//", "/"),
                "sample": environment.samples[metadata_source].sample,
                "atlas_pixel_size": atlas_pixel_size,
            }
        else:
            dcg_data = {
                "experiment_type_id": experiment_type_id,
                "tag": dcg_tag,
                "sample": environment.samples[metadata_source].sample,
            }
    capture_post(
        base_url=str(environment.url.geturl()),
        router_name="workflow.router",
        function_name="register_dc_group",
        token=token,
        instrument_name=environment.instrument_name,
        visit_name=environment.visit,
        session_id=environment.murfey_session,
        data=dcg_data,
    )
    return dcg_tag


def detect_acquisition_software(dir_for_transfer: Path) -> str:
    glob = dir_for_transfer.glob("*")
    for f in glob:
        if f.name.startswith("EPU") or f.name.startswith("GridSquare"):
            return "epu"
        if f.name.startswith("Position") or f.suffix == ".mdoc":
            return "tomo"
    return ""


class ProcessingParameter(NamedTuple):
    name: str
    label: str
    default: Any = None


class Context:
    user_params: list[ProcessingParameter] = []
    metadata_params: list[ProcessingParameter] = []

    def __init__(self, name: str, acquisition_software: str, token: str):
        self._acquisition_software = acquisition_software
        self._token = token
        self.name = name
        self.data_collection_parameters: dict = {}

    def post_transfer(
        self,
        transferred_file: Path,
        environment: MurfeyInstanceEnvironment | None = None,
        **kwargs,
    ):
        # Search external packages for additional hooks to include in Murfey
        for h in entry_points(group="murfey.post_transfer_hooks"):
            if h.name == self.name:
                h.load()(transferred_file, environment=environment, **kwargs)

    def post_first_transfer(
        self,
        transferred_file: Path,
        environment: MurfeyInstanceEnvironment | None = None,
        **kwargs,
    ):
        self.post_transfer(transferred_file, environment=environment, **kwargs)

    def gather_metadata(
        self, metadata_file: Path, environment: MurfeyInstanceEnvironment | None = None
    ) -> OrderedDict | None:
        raise NotImplementedError(
            f"gather_metadata must be declared in derived class to be used: {self}"
        )