"""
XERV CRAYON C-Extensions Package
================================

This package contains the native C/C++/CUDA extensions:

- crayon_cpu: AVX2/AVX-512 accelerated CPU tokenizer (always available)
- crayon_cuda: NVIDIA CUDA GPU tokenizer (optional, requires nvcc)
- crayon_rocm: AMD ROCm GPU tokenizer (optional, requires hipcc)

Import Behavior:
    - crayon_cpu is imported eagerly and will raise ImportError if missing
    - crayon_cuda and crayon_rocm are lazy-loaded to avoid import errors
    - Use check_* functions to safely probe availability

Example:
    >>> from crayon.c_ext import crayon_cpu
    >>> from crayon.c_ext import is_cuda_available, is_rocm_available
    >>> 
    >>> if is_cuda_available():
    ...     from crayon.c_ext import crayon_cuda
"""

import sys
from typing import Optional, Tuple

# ============================================================================
# CPU BACKEND (Required)
# ============================================================================

try:
    from . import crayon_cpu
except ImportError as e:
    # Provide helpful error message for common issues
    _cpu_error = (
        "Failed to import crayon_cpu extension. This is required for Crayon to work.\n"
        "Possible causes:\n"
        "  1. The package was not installed correctly (try: pip install --force-reinstall xerv-crayon)\n"
        "  2. The C++ extension failed to compile (check for compiler errors during install)\n"
        "  3. Python version mismatch (Crayon requires Python 3.10+)\n"
        f"Original error: {e}"
    )
    raise ImportError(_cpu_error) from e


# ============================================================================
# GPU BACKENDS (Optional - Lazy Import)
# ============================================================================

_cuda_module: Optional[object] = None
_rocm_module: Optional[object] = None
_cuda_checked: bool = False
_rocm_checked: bool = False
_cuda_error: Optional[str] = None
_rocm_error: Optional[str] = None


def is_cuda_available() -> bool:
    """
    Check if the CUDA backend is available.
    
    Returns:
        True if crayon_cuda can be imported and CUDA is functional.
    """
    global _cuda_checked, _cuda_module, _cuda_error
    
    if _cuda_checked:
        return _cuda_module is not None
    
    _cuda_checked = True
    try:
        from . import crayon_cuda as _cuda
        # Verify it's functional
        _ = _cuda.get_hardware_info()
        _cuda_module = _cuda
        return True
    except ImportError as e:
        _cuda_error = f"ImportError: {e}"
        return False
    except Exception as e:
        _cuda_error = f"RuntimeError: {e}"
        return False


def is_rocm_available() -> bool:
    """
    Check if the ROCm backend is available.
    
    Returns:
        True if crayon_rocm can be imported and ROCm is functional.
    """
    global _rocm_checked, _rocm_module, _rocm_error
    
    if _rocm_checked:
        return _rocm_module is not None
    
    _rocm_checked = True
    try:
        from . import crayon_rocm as _rocm
        # Verify it's functional
        info = _rocm.get_hardware_info()
        if isinstance(info, str) and "Device Not Found" in info:
            _rocm_error = info
            return False
        _rocm_module = _rocm
        return True
    except ImportError as e:
        _rocm_error = f"ImportError: {e}"
        return False
    except Exception as e:
        _rocm_error = f"RuntimeError: {e}"
        return False


def get_cuda_error() -> Optional[str]:
    """Get the error message if CUDA is unavailable."""
    is_cuda_available()  # Ensure check has run
    return _cuda_error


def get_rocm_error() -> Optional[str]:
    """Get the error message if ROCm is unavailable."""
    is_rocm_available()  # Ensure check has run
    return _rocm_error


def get_available_backends() -> Tuple[str, ...]:
    """
    Get list of available backends.
    
    Returns:
        Tuple of available backend names ("cpu", "cuda", "rocm").
    """
    backends = ["cpu"]
    if is_cuda_available():
        backends.append("cuda")
    if is_rocm_available():
        backends.append("rocm")
    return tuple(backends)


def get_backend_info() -> dict:
    """
    Get detailed information about all backends.
    
    Returns:
        Dictionary with backend status and hardware info.
    """
    info = {
        "cpu": {
            "available": True,
            "hardware": crayon_cpu.get_hardware_info() if hasattr(crayon_cpu, 'get_hardware_info') else "Unknown"
        }
    }
    
    if is_cuda_available():
        try:
            from . import crayon_cuda
            hw = crayon_cuda.get_hardware_info()
            info["cuda"] = {"available": True, "hardware": hw}
        except Exception as e:
            info["cuda"] = {"available": False, "error": str(e)}
    else:
        info["cuda"] = {"available": False, "error": _cuda_error}
    
    if is_rocm_available():
        try:
            from . import crayon_rocm
            hw = crayon_rocm.get_hardware_info()
            info["rocm"] = {"available": True, "hardware": hw}
        except Exception as e:
            info["rocm"] = {"available": False, "error": str(e)}
    else:
        info["rocm"] = {"available": False, "error": _rocm_error}
    
    return info


# ============================================================================
# CONDITIONAL IMPORTS FOR TYPE CHECKING
# ============================================================================

# These will fail at runtime if not available, which is intentional
# Use is_cuda_available() / is_rocm_available() before importing

__all__ = [
    "crayon_cpu",
    "is_cuda_available",
    "is_rocm_available",
    "get_cuda_error",
    "get_rocm_error",
    "get_available_backends",
    "get_backend_info",
]
