"""Shared runtime adapters for classification evaluation flows."""

from __future__ import annotations

import zipfile
from pathlib import Path
from typing import Any

import numpy as np
import torch

from .eval_model_loader import load_model_with_runtime


def setup_tensorrt() -> tuple[Any, Any]:
    """Import TensorRT and PyCUDA lazily to keep base runtime lightweight.

    Args:
        None.

    Returns:
        Tuple of imported ``(cuda, trt)`` modules.
    """
    from ultralytics.utils.checks import check_requirements  # type: ignore[reportMissingImports]

    check_requirements("pycuda")
    check_requirements("tensorrt==10.4.0")
    import pycuda.driver as cuda  # type: ignore[import-not-found]
    import tensorrt as trt  # type: ignore[import-not-found]

    return cuda, trt


def allocate_buffers(engine: Any, *, cuda: Any, trt: Any) -> tuple[Any, Any, Any, Any, Any]:
    """Create host and device buffers for TensorRT engine input/output.

    Args:
        engine: TensorRT engine object exposing tensor shapes.
        cuda: Imported PyCUDA driver module.
        trt: Imported TensorRT module.

    Returns:
        Tuple ``(h_input, d_input, h_output, d_output, stream)``.
    """
    h_input = cuda.pagelocked_empty(trt.volume(engine.get_tensor_shape("images")), dtype=np.float32)
    h_output = cuda.pagelocked_empty(trt.volume(engine.get_tensor_shape("output")), dtype=np.float32)
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    stream = cuda.Stream()
    return h_input, d_input, h_output, d_output, stream


def do_inference(
    _engine: Any,
    context: Any,
    h_input: Any,
    d_input: Any,
    h_output: Any,
    d_output: Any,
    stream: Any,
    *,
    cuda: Any,
) -> Any:
    """Run one TensorRT inference step and return host output buffer.

    Args:
        _engine: Unused TensorRT engine placeholder for API compatibility.
        context: TensorRT execution context.
        h_input: Host input buffer.
        d_input: Device input buffer.
        h_output: Host output buffer.
        d_output: Device output buffer.
        stream: CUDA stream used for async copies and execution.
        cuda: Imported PyCUDA driver module.

    Returns:
        Host output buffer containing inference results.
    """
    cuda.memcpy_htod_async(d_input, h_input, stream)
    context.execute_v2(bindings=[int(d_input), int(d_output)])
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    stream.synchronize()
    return h_output


def load_engine(trt_filename: str, *, trt: Any) -> tuple[Any, Any]:
    """Deserialize TensorRT engine file and build execution context.

    Args:
        trt_filename: Path to serialized TensorRT engine file.
        trt: Imported TensorRT module.

    Returns:
        Tuple of deserialized ``(engine, context)`` objects.
    """
    with Path(trt_filename).open("rb") as f:
        engine_data = f.read()
    runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
    engine = runtime.deserialize_cuda_engine(engine_data)
    if not engine:
        raise RuntimeError("Failed to deserialize TensorRT engine.")
    context = engine.create_execution_context()
    return engine, context


def load_classification_model(action_tracker: Any, runtime_framework: str = "pytorch") -> Any:
    """Load classification model artifact for the requested runtime.

    Args:
        action_tracker: Tracker-like object with model download methods.
        runtime_framework: Runtime descriptor used to select loader strategy.

    Returns:
        Runtime-specific loaded model artifact.
    """

    def _load_onnx(tracker: Any) -> Any:
        """Download ONNX artifact and initialize an ONNX Runtime session.

        Args:
            tracker: Tracker-like object used to download model artifacts.

        Returns:
            Initialized ONNX Runtime inference session.
        """
        import onnxruntime as ort  # type: ignore[import-not-found]

        tracker.download_model("model.onnx", model_type="exported")
        return ort.InferenceSession("model.onnx", None)

    def _load_openvino(tracker: Any) -> Any:
        """Download OpenVINO archive, compile model, and return output handle.

        Args:
            tracker: Tracker-like object used to download model artifacts.

        Returns:
            Tuple of ``(compiled_model, output_layer)``.
        """
        import openvino.runtime as ov  # type: ignore[import-not-found]

        tracker.download_model("model_openvino.zip", model_type="exported")
        with zipfile.ZipFile("model_openvino.zip", "r") as zip_ref:
            zip_ref.extractall("model_openvino")
        model_xml_path = "model_openvino/model.xml"
        core = ov.Core()
        device_name = "GPU" if "GPU" in core.available_devices else "CPU"
        ov_model = core.read_model(model=model_xml_path)
        compiled_model = core.compile_model(model=ov_model, device_name=device_name)
        output_layer = compiled_model.output(0)
        return compiled_model, output_layer

    def _load_torchscript(tracker: Any) -> Any:
        """Download TorchScript model and move it to active inference device.

        Args:
            tracker: Tracker-like object used to download model artifacts.

        Returns:
            Eval-ready TorchScript module.
        """
        tracker.download_model("model.torchscript", model_type="exported")
        model = torch.jit.load("model.torchscript", map_location="cpu")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device)
        model.eval()
        return model

    def _load_pytorch(tracker: Any) -> Any:
        """Download PyTorch checkpoint and prepare an eval-ready module.

        Args:
            tracker: Tracker-like object used to download model artifacts.

        Returns:
            Eval-ready PyTorch module.
        """
        tracker.download_model("model.pt")
        model = torch.load("model.pt", map_location="cpu", weights_only=False)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device)
        model.eval()
        return model

    def _load_tensorrt(tracker: Any) -> Any:
        """Download TensorRT engine and deserialize runtime engine/context.

        Args:
            tracker: Tracker-like object used to download model artifacts.

        Returns:
            Tuple of deserialized ``(engine, context)`` objects.
        """
        tracker.download_model("model.engine", model_type="exported")
        _, trt = setup_tensorrt()
        return load_engine("model.engine", trt=trt)

    return load_model_with_runtime(
        action_tracker=action_tracker,
        runtime_framework=runtime_framework,
        loaders={
            "pytorch": _load_pytorch,
            "torchscript": _load_torchscript,
            "onnx": _load_onnx,
            "openvino": _load_openvino,
            "tensorrt": _load_tensorrt,
        },
        default_key="pytorch",
    )


def get_tensorrt_inference_results(loader: Any, model: tuple[Any, Any], action_tracker: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Run classification inference for all batches on TensorRT runtime.

    Args:
        loader: Classification dataloader yielding ``(images, target)``.
        model: Tuple of TensorRT ``(engine, context)`` objects.
        action_tracker: Tracker-like object (unused in TensorRT path).

    Returns:
        Tuple of concatenated ``(predictions, outputs, targets)`` tensors.
    """
    del action_tracker
    engine, context = model
    cuda, trt = setup_tensorrt()
    h_input, d_input, h_output, d_output, stream = allocate_buffers(engine, cuda=cuda, trt=trt)

    all_outputs = []
    all_targets = []
    all_predictions = []

    for images, target in loader:
        images_np = images.numpy().astype(np.float32)
        h_input = np.ravel(images_np)
        output = do_inference(engine, context, h_input, d_input, h_output, d_output, stream, cuda=cuda)
        output_tensor = torch.as_tensor(output)
        if output_tensor.ndim == 1:
            output_tensor = output_tensor.unsqueeze(0)
        predictions = torch.argmax(output_tensor, dim=1)
        all_predictions.append(predictions)
        all_outputs.append(output_tensor)
        all_targets.append(target)

    return (
        torch.cat(all_predictions, dim=0),
        torch.cat(all_outputs, dim=0),
        torch.cat(all_targets, dim=0),
    )


def get_onnx_inference_results(loader: Any, model: Any, action_tracker: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Run classification inference for all batches on ONNX Runtime.

    Args:
        loader: Classification dataloader yielding ``(images, target)``.
        model: ONNX Runtime session object.
        action_tracker: Tracker-like object (unused in ONNX path).

    Returns:
        Tuple of concatenated ``(predictions, outputs, targets)`` tensors.
    """
    del action_tracker
    all_outputs = []
    all_targets = []
    all_predictions = []

    for images, target in loader:
        input_data = {"images": images.numpy()}
        output = model.run(None, input_data)
        predictions = torch.argmax(torch.tensor(output[0]), dim=1)
        all_predictions.append(predictions)
        all_outputs.append(torch.tensor(output[0]))
        all_targets.append(target)

    return torch.cat(all_predictions, dim=0), torch.cat(all_outputs, dim=0), torch.cat(all_targets, dim=0)


def get_openvino_inference_results(loader: Any, model: tuple[Any, Any], action_tracker: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Run classification inference for all batches on OpenVINO runtime.

    Args:
        loader: Classification dataloader yielding ``(images, target)``.
        model: Tuple of ``(compiled_model, output_layer)`` objects.
        action_tracker: Tracker-like object (unused in OpenVINO path).

    Returns:
        Tuple of concatenated ``(predictions, outputs, targets)`` tensors.
    """
    del action_tracker
    compiled_model, output_layer = model
    all_outputs = []
    all_targets = []
    all_predictions = []

    for images, target in loader:
        results = compiled_model(images.numpy())
        output = torch.tensor(results[output_layer])
        predictions = np.argmax(output, axis=1)
        all_predictions.append(torch.tensor(predictions))
        all_outputs.append(torch.tensor(output))
        all_targets.append(target)

    return torch.cat(all_predictions, dim=0), torch.cat(all_outputs, dim=0), torch.cat(all_targets, dim=0)


def validate_device(loaded_model: torch.nn.Module) -> tuple[torch.nn.Module, torch.device]:
    """Move model to best available device with safe CPU fallback.

    Args:
        loaded_model: PyTorch module to validate on target device.

    Returns:
        Tuple of ``(model_on_device, resolved_device)``.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    try:
        loaded_model = loaded_model.to(device)
        fake_image = torch.randn(1, 3, 224, 224).to(device)
        with torch.no_grad():
            _ = loaded_model(fake_image)
    except Exception:
        device = torch.device("cpu")
        loaded_model = loaded_model.to(device)
    return loaded_model, device


def get_pytorch_inference_results(loader: Any, model: Any, action_tracker: Any) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Run classification inference for all batches on PyTorch runtime.

    Args:
        loader: Classification dataloader yielding ``(images, target)``.
        model: PyTorch or TorchScript module.
        action_tracker: Tracker-like object for optional inference logging.

    Returns:
        Tuple of concatenated ``(predictions, outputs, targets)`` tensors.
    """
    model, device = validate_device(model)
    model.eval()
    all_outputs = []
    all_targets = []
    all_predictions = []

    with torch.no_grad():
        for images, target in loader:
            images = images.to(device)
            target = target.to(device)
            output = model(images)
            predictions = torch.argmax(output, dim=1)

            if hasattr(action_tracker, "store_inference_results"):
                action_tracker.store_inference_results(
                    images=images,
                    outputs=output,
                    targets=target,
                    split_type="val",
                    project_type="classification",
                )

            all_predictions.append(predictions)
            all_outputs.append(output)
            all_targets.append(target)

    return torch.cat(all_predictions, dim=0), torch.cat(all_outputs, dim=0), torch.cat(all_targets, dim=0)

