"""
Clean CuPy Registry Implementation

Implements clean abstraction with internal library-specific logic.
All CuPy-specific details (GPU handling, CuCIM integration, etc.)
are handled internally without leaking into the ABC.
"""
from __future__ import annotations

import numpy as np
from typing import Tuple, List

from openhcs.constants import MemoryType
from openhcs.core.utils import optional_import
from .unified_registry import LibraryRegistryBase, RuntimeTestingRegistryBase

cp = optional_import("cupy")
cucim = optional_import("cucim")
cucim_skimage = optional_import("cucim.skimage")


class CupyRegistry(RuntimeTestingRegistryBase):
    """Clean CuPy registry with internal GPU handling logic."""

    # Library-specific exclusions (uses common ones)
    EXCLUSIONS = LibraryRegistryBase.COMMON_EXCLUSIONS

    # Modules to scan for functions
    MODULES_TO_SCAN = ['filters', 'morphology', 'measure', 'segmentation',
                       'feature', 'restoration', 'transform', 'exposure',
                       'color', 'util']

    # Memory type for this registry
    MEMORY_TYPE = MemoryType.CUPY.value

    # Float dtype for this registry
    FLOAT_DTYPE = cp.float32

    def __init__(self):
        super().__init__("cupy")

    # ===== ESSENTIAL ABC METHODS =====
    def get_library_version(self) -> str:
        return cucim.__version__

    def is_library_available(self) -> bool:
        return bool(cp) and bool(cucim_skimage)

    def get_library_object(self):
        return cucim_skimage

    def get_module_patterns(self) -> List[str]:
        """Get module patterns for CuPy (includes cucim patterns)."""
        return ['cupy', 'cucim']

    def get_display_name(self) -> str:
        """Get proper display name for CuPy."""
        return 'CuPy'

    # ===== HOOK IMPLEMENTATIONS =====
    def _create_array(self, shape: Tuple[int, ...], dtype):
        try:
            return cp.random.rand(*shape).astype(dtype)
        except Exception as e:
            # If CUDA initialization fails, raise a more descriptive error
            raise RuntimeError(f"CUDA initialization failed during CuPy array creation: {e}") from e

    def _check_first_parameter(self, first_param, func_name: str) -> bool:
        return first_param.name.lower() in {'image', 'input', 'array', 'img'}

    def _preprocess_input(self, image, func_name: str):
        return image  # No preprocessing needed for CuPy

    def _postprocess_output(self, result, original_image, func_name: str):
        # ProcessingContract system handles dimensional behavior - no categorization needed
        return result

    # ===== LIBRARY-SPECIFIC IMPLEMENTATIONS =====
    def _generate_function_name(self, name: str, module_name: str) -> str:
        """Generate function name - original for filters, prefixed for others."""
        return name if module_name == 'filters' else f"{module_name}_{name}"

    def _generate_tags(self, func_name: str) -> List[str]:
        """Generate tags with GPU tag."""
        tags = func_name.lower().replace("_", " ").split()
        tags.append("gpu")
        return tags

    def _stack_2d_results(self, func, test_3d):
        """Stack 2D results using CuPy."""
        results = [func(test_3d[z]) for z in range(test_3d.shape[0])]
        return cp.stack(results)

    def _arrays_close(self, arr1, arr2):
        """Compare arrays using CuPy."""
        return np.allclose(arr1.get(), arr2.get(), rtol=1e-5, atol=1e-8)


