import os
import sys
from pathlib import Path
from typing import Any, cast

# Direct interpreter launches do not apply Pixi's native-library PATH activation.
native_dll_dir = Path(sys.prefix) / "Library" / "bin"
if os.name == "nt" and native_dll_dir.is_dir():
    os.environ["PATH"] = str(native_dll_dir) + os.pathsep + os.environ.get("PATH", "")

import ndv
import xarray as xr
from czitools.read_tools import read_tools
from czitools.utils.napari_tools import display_xarray_in_napari
from czitools.utils.ndv_tools import _create_luts_ndv
from magicgui import magicgui
from magicgui.types import FileDialogMode

# Make the shared utility importable from any workshop subdirectory.
for utility_dir in (Path.cwd().resolve(), *Path.cwd().resolve().parents):
    if (utility_dir / "workshop_utils.py").is_file():
        if str(utility_dir) not in sys.path:
            sys.path.insert(0, str(utility_dir))
        break
else:
    raise FileNotFoundError("Start from the workshop root or a subdirectory.")

from workshop_utils import initialize_workshop

workshop_root, project_dir, data_dir, model_dir = initialize_workshop("czi_handling")


show_napari = False
show_ndv = True
use_mdstack = True
# option to toggle using a file open dialog or a hardcoded filename
use_dialog = True

# option to show the link to the dash dashboard when using dask (only relevant if use_mdstack=True)
show_dask = False

# Keep shared variables initialized for both code paths.
array6d: Any = None
subset_planes: dict[str, Any] = {}
result: Any = None


# open simple dialog to select a CZI file
@magicgui(
    filepath={
        "label": "Choose CZI files:",
        "mode": FileDialogMode.EXISTING_FILE,
        "filter": "*.czi",
        "value": data_dir,
    },
    call_button="Open CZI File",
)
def filespicker(filepath: Path) -> Path:
    """Take a filename and do something with it."""
    # Close the dialog after the file is selected and the button was pressed
    # This will also return the filepath to the caller
    filespicker.close()

    return filepath


if use_dialog:
    filespicker.filepath.changed.connect(print)
    filespicker.show(run=True)

    filepath = str(filespicker.filepath.value)
    print(f"Selected file: {filepath}")

elif not use_dialog:
    filepath = str(data_dir / "CellDivision_T10_Z15_CH2_DCV_small.czi")


if not use_mdstack:
    # return a xarray with dimension order STCZYX(A)
    # when no planes are specified the complete dataset
    # when planes are specified the metadata will be adapted accordingly for SizeS, SizeT, SizeC and SizeZ
    array6d, mdata = read_tools.read_6darray(
        filepath,
        zoom=1.0,
        # planes={"S": (0, 0), "T": (0, 2), "C": (0, 0), "Z": (0, 4)},
        adapt_metadata=True,  # adapt metadata according to the planes specified (SizeS, SizeT, SizeC, SizeZ)
        use_dask=False,  # use lazy dask arrays (data read only when accessed)
        use_xarray=True,  # return xarray DataArray with labeled dimensions (STCZYX(A))
    )

    if array6d is None:
        raise RuntimeError("read_6darray returned no data.")

    # get the planes
    attrs = getattr(array6d, "attrs", {})
    subset_planes = attrs.get("subset_planes", {}) if isinstance(attrs, dict) else {}

    # print the shape of the xarray etc.
    print(f"Shape: {array6d.shape}")
    print(f"Dimensions: {array6d.dims}")
    print(f"Subset Planes: {subset_planes}")

    if isinstance(attrs, dict):
        for k, v in attrs.items():
            print(f"{k} :  {v}")

elif use_mdstack:
    result, dims, num_stacks, mdata = read_tools.read_stacks(
        filepath,
        zoom=1.0,
        # planes={"S": (0, 0), "T": (0, 30), "C": (0, 0)},
        use_dask=True,  # use lazy dask arrays (data read only when accessed)
        use_xarray=True,  # return xarray DataArray with labeled dimensions (STCZYX(A))
        stack_scenes=True,  # stack scenes with the same shape into one array (if False, return a list of arrays)
        adapt_metadata=True,  # adapt metadata according to the planes specified (SizeS, SizeT, SizeC, SizeZ)
    )

    from dask.distributed import Client, LocalCluster

    if show_dask:
        # On Windows, process workers re-import this module and can re-open the file dialog.
        # Use threads in a single process to keep dashboard support without re-running top-level code.
        cluster = LocalCluster(processes=False)
        client = Client(cluster)
        print(f"Link to Dask dashboard: {client.dashboard_link}")

    print("\n=== Results ===")
    print(f"Number of stacks: {num_stacks}")
    print(f"Dimension order: {dims}")

    if isinstance(result, list):
        # List of per-scene arrays
        for idx, arr in enumerate(result):
            if isinstance(arr, xr.DataArray):
                print(f"Stack {idx}: dims={arr.dims}, shape={arr.shape}, dtype={arr.dtype}")
            else:
                print(f"Stack {idx}: shape={arr.shape}, dims={dims}, dtype={arr.dtype}")
    else:
        # Single stacked array
        if isinstance(result, xr.DataArray):
            print(f"Stacked: dims={result.dims}, shape={result.shape}, dtype={result.dtype}")
        else:
            print(f"Stacked: shape={result.shape}, dims={dims}, dtype={result.dtype}")

        # With use_dask=True, result is backed by dask - no data loaded yet
        print(f"\nArray shape (no data loaded): {result.shape}")

###  SHOW RESULTS IN VIEWER

if show_napari and array6d is not None:
    # Delegate Napari display to the utility function
    display_xarray_in_napari(array6d, mdata, subset_planes)

if show_ndv:
    # Determine the appropriate data to display in NDV based on the results from read_stacks
    viewer_data = None
    if isinstance(result, list):
        viewer_data = result[0] if len(result) > 0 else None
    elif result is not None:
        viewer_data = result
    elif array6d is not None:
        viewer_data = array6d

    if viewer_data is None:
        raise RuntimeError("No data available for NDV display.")

    # create luts for NDV based on the metadata (e.g. channel names, colors, etc.)
    luts: Any = _create_luts_ndv(cast(Any, mdata))

    # display in NDV (viewer_data can be a single xarray DataArray or a list of DataArrays)
    viewer = ndv.imshow(viewer_data, channel_mode="composite", luts=luts)