"""Backend selection for the supported CPU implementation."""

from __future__ import annotations

from importlib.util import find_spec
from typing import TYPE_CHECKING, Protocol

from ._core.accelerated.compact import (
    ArrowCompactFacts,
    ColumnArticleBatch,
    CompactColumnScreenBatch,
    CompactRelevanceIndicator,
)
from .models import BackendInfo

if TYPE_CHECKING:
    from .customization import Taxonomy


def _native_available() -> bool:
    try:
        return find_spec("range_curate._native") is not None
    except (ImportError, ModuleNotFoundError, ValueError):
        return False


class _Backend(Protocol):
    name: str

    def describe(
        self,
        indicator: CompactRelevanceIndicator,
    ) -> BackendInfo: ...

    def screen_batch(
        self,
        indicator: CompactRelevanceIndicator,
        batch: ColumnArticleBatch,
        *,
        workers: int,
    ) -> CompactColumnScreenBatch: ...

    def screen_arrow(
        self,
        indicator: CompactRelevanceIndicator,
        taxonomy: "Taxonomy",
        data,
        *,
        text_column: str,
        timestamp_column: str,
        workers: int,
        row_count: int | None,
    ) -> ArrowCompactFacts: ...


class _CpuBackend:
    name = "cpu"

    def describe(
        self,
        indicator: CompactRelevanceIndicator,
    ) -> BackendInfo:
        native_list = indicator.list_available
        native_arrow = indicator.arrow_available
        if native_list and native_arrow:
            detail = "Rust list and Arrow screening are available"
        elif native_list:
            detail = "Rust list screening is available; Arrow screening is unavailable"
        else:
            detail = (
                "Python prediction is available; rebuild the native extension "
                "for compact screening"
            )
        return BackendInfo(
            name=self.name,
            available=True,
            accelerated=native_list,
            detail=detail,
        )

    def screen_batch(
        self,
        indicator: CompactRelevanceIndicator,
        batch: ColumnArticleBatch,
        *,
        workers: int,
    ) -> CompactColumnScreenBatch:
        if not indicator.list_available:
            raise RuntimeError(
                "Compact screening requires the packaged Rust extension; "
                "install a wheel for this platform or rebuild the package"
            )
        return indicator.screen_column_batch(batch, workers=workers)

    def screen_arrow(
        self,
        indicator: CompactRelevanceIndicator,
        taxonomy: "Taxonomy",
        data,
        *,
        text_column: str,
        timestamp_column: str,
        workers: int,
        row_count: int | None,
    ) -> ArrowCompactFacts:
        del taxonomy
        if not indicator.arrow_available:
            raise RuntimeError(
                "Arrow screening requires the packaged Rust extension and the "
                "'arrow' extra"
            )
        return indicator.screen_arrow(
            data,
            text_column=text_column,
            timestamp_column=timestamp_column,
            workers=workers,
            row_count=row_count,
        )


class _CudaBackend(_CpuBackend):
    """CUDA bulk Arrow screening; rich/list operations remain CPU work."""
    name = "cuda"

    def __init__(self):
        import threading
        if find_spec("range_curate_cuda") is None:
            raise RuntimeError("Install the matching range-curate-cuda wheel for CUDA Arrow screening")
        self._engine = None
        self._lock = threading.Lock()

    def describe(self, indicator):
        return BackendInfo(name="cuda", available=True, accelerated=True,
            detail="Phase60 CUDA for bulk Arrow screening; rich/list operations use CPU. Device support is validated on first Arrow call.")

    def screen_arrow(self, indicator, taxonomy, data, *, text_column, timestamp_column, workers, row_count):
        if self._engine is None:
            with self._lock:
                if self._engine is None:
                    from range_curate_cuda import CudaRelevanceEngine
                    self._engine = CudaRelevanceEngine(taxonomy=taxonomy, workers=workers)
        return self._engine.screen_arrow(data, text_column=text_column,
            timestamp_column=timestamp_column, workers=workers, row_count=row_count)


def create_backend(name: str) -> _Backend:
    normalized = name.strip().lower() if isinstance(name, str) else ""
    if normalized == "cpu":
        return _CpuBackend()
    if normalized == "cuda":
        return _CudaBackend()
    raise ValueError("backend must be 'cpu' or 'cuda'")


def available_backends() -> tuple[BackendInfo, ...]:
    """Discover packages without importing CuPy or initializing CUDA."""
    native = _native_available()
    cuda = find_spec("range_curate_cuda") is not None
    return (
        BackendInfo(name="cpu", available=True, accelerated=native,
            detail="Native compact screening" if native else "Python rich predictions; native screening unavailable"),
        BackendInfo(name="cuda", available=cuda, accelerated=cuda,
            detail="Optional Phase60 Arrow runtime installed; device support checked on use" if cuda
            else "Install the matching range-curate-cuda wheel; Windows CPython 3.10 and sm75 only"),
    )


__all__ = ["available_backends", "create_backend"]