"""
Video preprocessors for BADAS inference.

Provides GPU-accelerated preprocessing components:
- ZeroCopyTensorRTEngine: DEPRECATED — superseded by TensorRTBackend in
  badas.inference.backends, which implements the full InferenceBackend protocol
  and fixes the missing sigmoid in the return value.  Kept for backwards
  compatibility only; emits DeprecationWarning on instantiation.
"""

import warnings

import torch

from badas.core.preprocessing import VideoPreprocessor  # noqa: F401 (re-exported)


class ZeroCopyTensorRTEngine:
    """
    DEPRECATED: Use ``TensorRTBackend`` from ``badas.inference.backends`` instead.

    This class is a legacy standalone TensorRT wrapper that predates the
    ``InferenceBackend`` protocol.  It has been superseded by
    ``TensorRTBackend``, which provides the same zero-copy functionality
    through a richer, protocol-conformant API.

    Relationship to TensorRTBackend
    --------------------------------
    ``ZeroCopyTensorRTEngine`` was the original zero-copy TRT helper.
    ``TensorRTBackend`` replaced it and added:

    - Full ``InferenceBackend`` protocol (``forward``, ``forward_async``,
      ``collect_result``, ``forward_batch``).
    - A dedicated CUDA stream (``_trt_stream``) that avoids the extra
      ``cudaStreamSynchronize`` calls TensorRT inserts on stream 0.
    - ``forward_zero_copy(tensor)`` — equivalent to ``infer_from_tensor``,
      but with proper stream ordering.
    - ``forward_from_ptr(ptr)`` — equivalent to ``infer_from_ptr``, with
      proper stream ordering.
    - Correct sigmoid application: both methods return a probability in
      [0, 1].  ``ZeroCopyTensorRTEngine.infer_from_tensor`` and
      ``infer_from_ptr`` return the raw logit via ``.item()`` without
      sigmoid, which was a latent correctness bug.

    Migration
    ---------
    Replace::

        engine = ZeroCopyTensorRTEngine("model.trt")
        prob = engine.infer_from_tensor(tensor)
        prob = engine.infer_from_ptr(ptr)

    With::

        from badas.inference.backends import TensorRTBackend
        backend = TensorRTBackend("model.trt")
        prob = backend.forward_zero_copy(tensor)   # was infer_from_tensor
        prob = backend.forward_from_ptr(ptr)        # was infer_from_ptr
        prob = backend.forward(tensor)              # standard (copies to internal buffer)

    Or simply use the high-level ``BADAS`` class, which selects the correct
    backend automatically based on the model file extension.
    """

    def __init__(
        self,
        engine_path: str,
        input_name: str = "pixel_values_videos",
        output_name: str = "logits",
        quiet: bool = False,
    ):
        warnings.warn(
            "ZeroCopyTensorRTEngine is deprecated and will be removed in a future release. "
            "Use TensorRTBackend from badas.inference.backends instead. "
            "See the ZeroCopyTensorRTEngine docstring for a migration guide.",
            DeprecationWarning,
            stacklevel=2,
        )

        try:
            import tensorrt as trt
        except ImportError:
            raise ImportError("tensorrt not installed. Install with: pip install tensorrt")

        self.quiet = quiet
        self.input_name = input_name
        self.output_name = output_name

        self._log(f"Loading TensorRT engine: {engine_path}")
        logger = trt.Logger(trt.Logger.WARNING)
        runtime = trt.Runtime(logger)

        with open(engine_path, 'rb') as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())

        if self.engine is None:
            raise RuntimeError(f"Failed to load TensorRT engine: {engine_path}")

        self.context = self.engine.create_execution_context()
        self._setup_io(trt)

    def _log(self, msg: str):
        if not self.quiet:
            print(f"   {msg}")

    def _setup_io(self, trt):
        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            shape = self.engine.get_tensor_shape(name)
            mode = self.engine.get_tensor_mode(name)

            if mode == trt.TensorIOMode.INPUT:
                self.input_name = name
                self.input_shape = tuple(shape)
                self._log(f"TRT Input: {name} {list(shape)}")
            else:
                self.output_name = name
                self.output_shape = tuple(shape) if shape[0] > 0 else (1,)
                self._log(f"TRT Output: {name} {list(shape)}")

        self.context.set_input_shape(self.input_name, self.input_shape)
        self.output_buffer = torch.empty(
            self.output_shape, dtype=torch.float32, device='cuda'
        ).contiguous()
        self._log("Zero-copy TensorRT engine ready")

    def infer_from_tensor(self, tensor: torch.Tensor) -> float:
        """
        Run zero-copy inference from a CUDA tensor.

        .. deprecated::
            Use ``TensorRTBackend.forward_zero_copy(tensor)`` instead, which
            runs on a dedicated CUDA stream and applies sigmoid so the return
            value is a true probability in [0, 1].

        Warning: this method returns the raw logit (pre-sigmoid), NOT a
        probability.  This is a known bug in the legacy API.

        Args:
            tensor: Preprocessed video tensor on CUDA, contiguous, FP32.

        Returns:
            Raw logit scalar (NOT sigmoid-transformed).
        """
        if not tensor.is_cuda:
            raise ValueError("Tensor must be on CUDA device")
        if not tensor.is_contiguous():
            tensor = tensor.contiguous()
        if tensor.dtype != torch.float32:
            tensor = tensor.float()

        self.context.set_tensor_address(self.input_name, tensor.data_ptr())
        self.context.set_tensor_address(self.output_name, self.output_buffer.data_ptr())
        self.context.execute_async_v3(torch.cuda.current_stream().cuda_stream)
        torch.cuda.synchronize()
        return float(self.output_buffer.item())

    def infer_from_ptr(self, input_ptr: int) -> float:
        """
        Run zero-copy inference from a raw GPU memory pointer.

        .. deprecated::
            Use ``TensorRTBackend.forward_from_ptr(ptr)`` instead, which
            runs on a dedicated CUDA stream and applies sigmoid so the return
            value is a true probability in [0, 1].

        Warning: this method returns the raw logit (pre-sigmoid), NOT a
        probability.  This is a known bug in the legacy API.

        Args:
            input_ptr: GPU memory pointer to a contiguous FP32 input tensor.

        Returns:
            Raw logit scalar (NOT sigmoid-transformed).
        """
        self.context.set_tensor_address(self.input_name, input_ptr)
        self.context.set_tensor_address(self.output_name, self.output_buffer.data_ptr())
        self.context.execute_async_v3(torch.cuda.current_stream().cuda_stream)
        torch.cuda.synchronize()
        return float(self.output_buffer.item())


__all__ = [
    'ZeroCopyTensorRTEngine',
]
