#!/usr/bin/env python3
"""
zc_bind.py - Python ctypes wrapper for libcdna_zc.so

Provides zero-copy memory mapping with GPU registration:
  - mmap file regions
  - mlock to pin pages
  - cudaHostRegister for GPU direct access
  - DLPack export for framework interop

Work Order: WO-ZC-SHIM

Usage:
    from helix_cdc.zc.zc_bind import map_register, unregister_unmap, HostSpan

    # Map a CDNA shard region
    span = map_register("/path/to/shard.cdna2.hxz", offset=1024, nbytes=4096)

    # Access data
    print(f"Host ptr: {span.host_ptr}")
    print(f"GPU alias: {span.dev_alias}")

    # Clean up
    unregister_unmap(span)
"""

from __future__ import annotations

import ctypes as C
from ctypes import c_void_p, c_char_p, c_int, c_int64, c_size_t, c_uint8
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

# Error codes matching cdna_zc.h
class ZcError:
    OK = 0
    ERR_OPEN = -1
    ERR_MMAP = -2
    ERR_MLOCK = -3
    ERR_CUDA_REG = -4
    ERR_CUDA_DEV = -5
    ERR_INVALID = -6
    ERR_MUNLOCK = -7
    ERR_MUNMAP = -8
    ERR_CUDA_UNREG = -9

    @classmethod
    def to_str(cls, code: int) -> str:
        mapping = {
            cls.OK: "OK",
            cls.ERR_OPEN: "Failed to open file",
            cls.ERR_MMAP: "mmap failed",
            cls.ERR_MLOCK: "mlock failed",
            cls.ERR_CUDA_REG: "cudaHostRegister failed",
            cls.ERR_CUDA_DEV: "cudaHostGetDevicePointer failed",
            cls.ERR_INVALID: "Invalid arguments",
            cls.ERR_MUNLOCK: "munlock failed",
            cls.ERR_MUNMAP: "munmap failed",
            cls.ERR_CUDA_UNREG: "cudaHostUnregister failed",
        }
        return mapping.get(code, f"Unknown error {code}")


class HostSpan(C.Structure):
    """
    Mapped memory region with GPU registration.

    Matches C struct:
        void* host_ptr;
        void* dev_alias;
        size_t nbytes;
        size_t page_size;
        int fd;
        int64_t offset;
        unsigned int flags;
    """
    _fields_ = [
        ("host_ptr", c_void_p),      # Host-side mmap'd pointer
        ("dev_alias", c_void_p),     # GPU-visible pointer
        ("nbytes", c_size_t),        # Size of mapped region
        ("page_size", c_size_t),     # System page size
        ("fd", c_int),               # File descriptor
        ("offset", c_int64),         # Offset within file
        ("flags", C.c_uint),         # Registration flags
    ]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for debugging/receipts."""
        return {
            "host_ptr": hex(self.host_ptr) if self.host_ptr else None,
            "dev_alias": hex(self.dev_alias) if self.dev_alias else None,
            "nbytes": self.nbytes,
            "page_size": self.page_size,
            "fd": self.fd,
            "offset": self.offset,
            "flags": self.flags,
        }


# DLPack data types
class DLDataType:
    kDLInt = 0
    kDLUInt = 1
    kDLFloat = 2
    kDLBfloat = 4


# Library path finding
def _find_library() -> Optional[Path]:
    """Find libcdna_zc.so in expected locations."""
    candidates = [
        Path(__file__).parent / "libcdna_zc.so",
        Path("/usr/local/lib/libcdna_zc.so"),
        Path("/usr/lib/libcdna_zc.so"),
    ]
    for p in candidates:
        if p.exists():
            return p
    return None


# Lazy library loading
_lib: Optional[C.CDLL] = None
_lib_path: Optional[Path] = None


def _get_lib() -> C.CDLL:
    """Get or load the C library."""
    global _lib, _lib_path

    if _lib is not None:
        return _lib

    lib_path = _find_library()
    if lib_path is None:
        raise RuntimeError(
            "libcdna_zc.so not found. Build it with:\n"
            "  cd helix_cdc/zc && make"
        )

    _lib = C.CDLL(str(lib_path))
    _lib_path = lib_path

    # Set up function signatures
    _lib.cdna_map_register.argtypes = [c_char_p, c_int64, c_size_t, C.POINTER(HostSpan), c_int]
    _lib.cdna_map_register.restype = c_int

    _lib.cdna_unregister_unmap.argtypes = [C.POINTER(HostSpan)]
    _lib.cdna_unregister_unmap.restype = c_int

    _lib.cdna_make_dlpack_cpu.argtypes = [c_void_p, C.POINTER(c_int64), c_int, c_uint8, c_uint8, C.POINTER(c_void_p)]
    _lib.cdna_make_dlpack_cpu.restype = c_int

    _lib.cdna_get_page_size.argtypes = []
    _lib.cdna_get_page_size.restype = c_size_t

    _lib.cdna_get_cuda_device_count.argtypes = []
    _lib.cdna_get_cuda_device_count.restype = c_int

    _lib.cdna_get_bar1_free_bytes.argtypes = [c_int]
    _lib.cdna_get_bar1_free_bytes.restype = c_size_t

    return _lib


def map_register(
    path: str,
    offset: int,
    nbytes: int,
    onfault: bool = True,
) -> HostSpan:
    """
    Map file region and register for GPU zero-copy access.

    Args:
        path: Path to file (CDNA shard, etc.)
        offset: Byte offset within file
        nbytes: Number of bytes to map
        onfault: If True, eagerly fault pages (MAP_POPULATE)

    Returns:
        HostSpan with mapped memory info

    Raises:
        OSError: On mapping/registration failure
    """
    lib = _get_lib()

    span = HostSpan()
    result = lib.cdna_map_register(
        path.encode("utf-8"),
        c_int64(offset),
        c_size_t(nbytes),
        C.byref(span),
        c_int(1 if onfault else 0),
    )

    if result != ZcError.OK:
        raise OSError(f"map_register failed: {ZcError.to_str(result)} (code {result})")

    return span


def unregister_unmap(span: HostSpan) -> None:
    """
    Unregister from GPU and unmap memory.

    Args:
        span: HostSpan to clean up

    Raises:
        OSError: On cleanup failure
    """
    lib = _get_lib()

    result = lib.cdna_unregister_unmap(C.byref(span))

    if result != ZcError.OK:
        raise OSError(f"unregister_unmap failed: {ZcError.to_str(result)} (code {result})")


def make_dlpack_cpu(
    host_ptr: int,
    shape: List[int],
    dtype: str = "float32",
) -> Any:
    """
    Create a DLPack capsule for the given host pointer.

    Args:
        host_ptr: Pointer to data (int address)
        shape: Tensor shape
        dtype: Data type ("float32", "uint8", etc.)

    Returns:
        PyCapsule containing DLManagedTensor
    """
    lib = _get_lib()

    # Map dtype to DLPack codes
    dtype_map = {
        "float32": (DLDataType.kDLFloat, 32),
        "float16": (DLDataType.kDLFloat, 16),
        "uint8": (DLDataType.kDLUInt, 8),
        "int32": (DLDataType.kDLInt, 32),
        "int64": (DLDataType.kDLInt, 64),
    }

    if dtype not in dtype_map:
        raise ValueError(f"Unsupported dtype: {dtype}")

    dtype_code, dtype_bits = dtype_map[dtype]

    # Create shape array
    ndim = len(shape)
    shape_arr = (c_int64 * ndim)(*shape)

    capsule_out = c_void_p()

    result = lib.cdna_make_dlpack_cpu(
        c_void_p(host_ptr),
        shape_arr,
        c_int(ndim),
        c_uint8(dtype_code),
        c_uint8(dtype_bits),
        C.byref(capsule_out),
    )

    if result != ZcError.OK:
        raise RuntimeError(f"make_dlpack_cpu failed: {ZcError.to_str(result)}")

    # Note: This returns raw managed tensor pointer
    # For framework interop, wrap in PyCapsule
    return capsule_out.value


def get_page_size() -> int:
    """Get system page size in bytes."""
    return _get_lib().cdna_get_page_size()


def get_cuda_device_count() -> int:
    """Get number of CUDA devices."""
    return _get_lib().cdna_get_cuda_device_count()


def get_bar1_free_bytes(device: int = 0) -> int:
    """Get free BAR1 memory on device."""
    return _get_lib().cdna_get_bar1_free_bytes(device)


def span_to_numpy(span: HostSpan, shape: Tuple[int, ...], dtype: np.dtype = np.uint8) -> np.ndarray:
    """
    Create a NumPy array view of the mapped memory.

    IMPORTANT: The span must remain valid for the lifetime of the array!

    Args:
        span: Mapped HostSpan
        shape: Desired array shape
        dtype: NumPy dtype

    Returns:
        NumPy array view (does not copy data)
    """
    if span.host_ptr is None or span.host_ptr == 0:
        raise ValueError("Invalid span: null host_ptr")

    expected_bytes = int(np.prod(shape)) * np.dtype(dtype).itemsize
    if expected_bytes > span.nbytes:
        raise ValueError(
            f"Shape {shape} with dtype {dtype} requires {expected_bytes} bytes, "
            f"but span only has {span.nbytes} bytes"
        )

    # Create array from pointer
    arr_type = C.c_uint8 * span.nbytes
    buffer = arr_type.from_address(span.host_ptr)

    # View as NumPy array
    arr = np.frombuffer(buffer, dtype=dtype, count=int(np.prod(shape)))
    return arr.reshape(shape)


# Context manager for easy cleanup
class MappedSpan:
    """
    Context manager for mapped memory regions.

    Usage:
        with MappedSpan("/path/to/file", offset=0, nbytes=4096) as span:
            arr = span_to_numpy(span.span, (64, 64), np.uint8)
            # Use array...
        # Automatically cleaned up
    """

    def __init__(self, path: str, offset: int, nbytes: int, onfault: bool = True):
        self.path = path
        self.offset = offset
        self.nbytes = nbytes
        self.onfault = onfault
        self.span: Optional[HostSpan] = None

    def __enter__(self) -> "MappedSpan":
        self.span = map_register(self.path, self.offset, self.nbytes, self.onfault)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        if self.span is not None:
            try:
                unregister_unmap(self.span)
            except OSError:
                pass  # Best effort cleanup
            self.span = None


def library_info() -> Dict[str, Any]:
    """Get library info for debugging."""
    return {
        "lib_path": str(_lib_path) if _lib_path else None,
        "page_size": get_page_size() if _lib else 0,
        "cuda_devices": get_cuda_device_count() if _lib else 0,
    }


if __name__ == "__main__":
    print("zc_bind.py - Zero-Copy CDNA Bindings")
    print()

    try:
        lib = _get_lib()
        print(f"Library loaded: {_lib_path}")
        print(f"Page size: {get_page_size()} bytes")
        print(f"CUDA devices: {get_cuda_device_count()}")

        if get_cuda_device_count() > 0:
            free_bytes = get_bar1_free_bytes(0)
            print(f"Device 0 free memory: {free_bytes / (1024**3):.2f} GB")
    except RuntimeError as e:
        print(f"ERROR: {e}")
        print()
        print("To build the library:")
        print("  cd helix_cdc/zc && make")
