#!/usr/bin/env python3
"""RDK X5 BPU policy runtime wrappers."""

from contextlib import contextmanager
from pathlib import Path
import ctypes
import os
import sys

import numpy as np


HERE = Path(__file__).parent.resolve()
DEFAULT_CPP_LIB = HERE / "cpp" / "libbpu_dnn_policy.so"


def flush_c_stdio():
    try:
        ctypes.CDLL(None).fflush(None)
    except Exception:
        pass


@contextmanager
def suppress_c_output(enabled=True):
    if not enabled:
        yield
        return
    sys.stdout.flush()
    sys.stderr.flush()
    flush_c_stdio()
    saved_fds = {fd: os.dup(fd) for fd in (1, 2)}
    devnull_fd = os.open(os.devnull, os.O_WRONLY)
    try:
        os.dup2(devnull_fd, 1)
        os.dup2(devnull_fd, 2)
        yield
    finally:
        flush_c_stdio()
        for fd, saved_fd in saved_fds.items():
            os.dup2(saved_fd, fd)
            os.close(saved_fd)
        os.close(devnull_fd)


def _first_bpu_core(bpu_cores):
    cores = tuple(int(core) for core in bpu_cores)
    if not cores:
        return 0
    core = cores[0]
    if core == 0:
        return 1
    if core == 1:
        return 2
    return 0


class BpuInferLibPolicy:
    """Default fast backend: direct C++ DNN API via ctypes."""

    backend_name = "cpp_dnn_api_x5"

    def __init__(self, model_path, priority=0, bpu_cores=(0,), cpp_lib=DEFAULT_CPP_LIB):
        self.model_path = Path(model_path).expanduser().resolve()
        self.cpp_lib = Path(cpp_lib).expanduser().resolve()
        if not self.model_path.exists():
            raise FileNotFoundError(f"BPU model not found: {self.model_path}")
        if not self.cpp_lib.exists():
            raise FileNotFoundError(
                f"C++ BPU runtime library not found: {self.cpp_lib}. "
                f"Build it on the RDK X5 board with: cd {self.cpp_lib.parent} && bash build_board.sh"
            )

        self.priority = int(priority)
        self.bpu_cores = tuple(int(core) for core in bpu_cores)
        self._lib = ctypes.CDLL(str(self.cpp_lib))
        self._lib.rlgym_bpu_create.argtypes = [
            ctypes.c_char_p,
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_char_p,
            ctypes.c_int,
        ]
        self._lib.rlgym_bpu_create.restype = ctypes.c_void_p
        self._lib.rlgym_bpu_infer.argtypes = [
            ctypes.c_void_p,
            ctypes.POINTER(ctypes.c_float),
            ctypes.POINTER(ctypes.c_float),
            ctypes.c_char_p,
            ctypes.c_int,
        ]
        self._lib.rlgym_bpu_infer.restype = ctypes.c_int
        self._lib.rlgym_bpu_destroy.argtypes = [ctypes.c_void_p]
        self._lib.rlgym_bpu_destroy.restype = None
        self._lib.rlgym_bpu_input_floats.argtypes = [ctypes.c_void_p]
        self._lib.rlgym_bpu_input_floats.restype = ctypes.c_int
        self._lib.rlgym_bpu_output_floats.argtypes = [ctypes.c_void_p]
        self._lib.rlgym_bpu_output_floats.restype = ctypes.c_int
        self._lib.rlgym_bpu_version.argtypes = []
        self._lib.rlgym_bpu_version.restype = ctypes.c_char_p

        err = ctypes.create_string_buffer(1024)
        self._handle = self._lib.rlgym_bpu_create(
            str(self.model_path).encode("utf-8"),
            _first_bpu_core(self.bpu_cores),
            self.priority,
            err,
            len(err),
        )
        if not self._handle:
            raise RuntimeError(err.value.decode("utf-8", errors="replace"))
        self.input_size = int(self._lib.rlgym_bpu_input_floats(self._handle))
        self.output_size = int(self._lib.rlgym_bpu_output_floats(self._handle))
        if self.input_size <= 0 or self.output_size <= 0:
            raise RuntimeError("invalid BPU tensor sizes reported by runtime")
        self.input_shape = (1, 1, 1, self.input_size)
        self.output_shape = (1, self.output_size, 1, 1)
        self._output = np.empty(self.output_size, dtype=np.float32)
        self._output_ptr = self._output.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
        self._err = ctypes.create_string_buffer(1024)

        print(f"[INFO] BPU model: {self.model_path}")
        print(f"[INFO]   Backend: {self.backend_name}")
        print(f"[INFO]   Runtime: {self.cpp_lib}")
        print(f"[INFO]   Input  : obs_4d {self.input_shape}")
        print(f"[INFO]   Output : actions {self.output_shape}")

    def close(self):
        handle = getattr(self, "_handle", None)
        if handle:
            self._lib.rlgym_bpu_destroy(handle)
            self._handle = None

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass

    def __call__(self, flat_input):
        arr = np.asarray(flat_input, dtype=np.float32)
        if arr.size != self.input_size:
            raise ValueError(f"BPU policy input has {arr.size} values, expected {self.input_size}")
        input_flat = np.ascontiguousarray(arr.reshape(-1), dtype=np.float32)
        rc = self._lib.rlgym_bpu_infer(
            self._handle,
            input_flat.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
            self._output_ptr,
            self._err,
            len(self._err),
        )
        if rc != 0:
            raise RuntimeError(self._err.value.decode("utf-8", errors="replace"))
        if not np.all(np.isfinite(self._output)):
            raise RuntimeError(f"BPU policy output is not finite: {self._output}")
        return self._output.copy()


class BpuInferLibPythonPolicy:
    """Slower Python package backend kept only for comparison."""

    backend_name = "bpu_infer_lib_x5_python"
    input_shape = (1, 1, 1, 450)
    output_shape = (1, 12, 1, 1)

    def __init__(self, model_path, priority=0, bpu_cores=(0,), suppress_runtime_output=False):
        self.model_path = Path(model_path).expanduser().resolve()
        if not self.model_path.exists():
            raise FileNotFoundError(f"BPU model not found: {self.model_path}")

        try:
            from bpu_infer_lib import Infer
        except ImportError as exc:
            raise RuntimeError(
                "bpu_infer_lib is required for the fallback Python BPU backend. "
                "Do not use hobot_dnn.pyeasy_dnn for this policy: it produced "
                "wrong actions in local board tests."
            ) from exc

        self.priority = int(priority)
        self.bpu_cores = tuple(int(core) for core in bpu_cores)
        self.suppress_runtime_output = bool(suppress_runtime_output)
        self.input_size = 450
        self.input_shape = (1, 1, 1, self.input_size)
        self.output_shape = (1, 12, 1, 1)
        with suppress_c_output(self.suppress_runtime_output):
            self.infer = Infer(False)
            loaded = self.infer.load_model(str(self.model_path))
        if not loaded:
            raise RuntimeError(f"Failed to load BPU model: {self.model_path}")

        print(f"[INFO] BPU model: {self.model_path}")
        print(f"[INFO]   Backend: {self.backend_name}")
        print(f"[INFO]   Input  : obs_4d {self.input_shape}")
        print(f"[INFO]   Output : actions {self.output_shape}")

    def __call__(self, flat_input):
        arr = np.asarray(flat_input, dtype=np.float32)
        if arr.size != self.input_size:
            raise ValueError(f"BPU policy input has {arr.size} values, expected {self.input_size}")
        input_4d = np.ascontiguousarray(arr.reshape(self.input_shape), dtype=np.float32)
        with suppress_c_output(self.suppress_runtime_output):
            copied = self.infer.read_numpy_arr_float32(input_4d, 0)
            if copied:
                self.infer.forward(True)
                got_output = self.infer.get_output()
                if not got_output:
                    raise RuntimeError("bpu_infer_lib get_output() failed")
                output = self.infer.get_infer_res_np_float32(0)
        if not copied:
            raise RuntimeError("Failed to copy float32 input into bpu_infer_lib")
        action = np.asarray(output, dtype=np.float32).reshape(-1)
        if action.size < 12:
            raise RuntimeError(f"BPU policy output has {action.size} values, expected at least 12")
        action = action[:12].astype(np.float32, copy=False)
        if not np.all(np.isfinite(action)):
            raise RuntimeError(f"BPU policy output is not finite: {action}")
        return action
