"""
QuarterBit PyTorch Backend — Utilities
Copyright (c) 2026 Clouthier Simulation Labs. All rights reserved.
"""

import ctypes
from pathlib import Path
import warnings
import sys

_lib = None
_available = None

def _load_backend():
    """Load the QuarterBit CUDA backend."""
    global _lib, _available

    if _available is not None:
        return _available

    # DLL search paths (compiled binaries only - source is confidential)
    candidates = [
        Path(__file__).parent.parent / "bin" / "quarterbit_core.dll",
        Path(__file__).parent.parent.parent / "build" / "quarterbit_core.dll",
        Path(r"C:\QuarterBit\build\quarterbit_core.dll"),
    ]

    # Linux paths
    if sys.platform != "win32":
        candidates = [
            Path(__file__).parent.parent / "bin" / "libquarterbit_core.so",
            Path(__file__).parent.parent.parent / "build" / "libquarterbit_core.so",
        ]

    for path in candidates:
        if path.exists():
            try:
                _lib = ctypes.CDLL(str(path))
                _configure_functions()
                _available = True
                return True
            except Exception as e:
                warnings.warn(f"Failed to load QuarterBit backend: {e}")

    _available = False
    return False

def _configure_functions():
    """Configure ctypes function signatures."""
    global _lib

    P_F = ctypes.POINTER(ctypes.c_float)

    # eft_accumulate
    _lib.eft_accumulate.argtypes = [P_F, P_F, P_F, ctypes.c_int]
    _lib.eft_accumulate.restype = ctypes.c_int

    # eft_accumulate_scaled
    _lib.eft_accumulate_scaled.argtypes = [P_F, P_F, P_F, ctypes.c_float, ctypes.c_int]
    _lib.eft_accumulate_scaled.restype = ctypes.c_int

    # eft_sgd_step
    _lib.eft_sgd_step.argtypes = [
        P_F, P_F, P_F,  # weight, comp, grad
        ctypes.c_float, ctypes.c_float,  # lr, weight_decay
        ctypes.c_int  # n
    ]
    _lib.eft_sgd_step.restype = ctypes.c_int

    # eft_adam_step
    _lib.eft_adam_step.argtypes = [
        P_F, P_F, P_F,  # weight, weight_comp, grad
        P_F, P_F,  # exp_avg, exp_avg_sq
        P_F, P_F,  # exp_avg_comp, exp_avg_sq_comp
        ctypes.c_float, ctypes.c_float, ctypes.c_float,  # lr, beta1, beta2
        ctypes.c_float, ctypes.c_float,  # eps, weight_decay
        ctypes.c_int, ctypes.c_int  # step, n
    ]
    _lib.eft_adam_step.restype = ctypes.c_int

    # eft_matmul
    _lib.eft_matmul.argtypes = [P_F, P_F, P_F, ctypes.c_int, ctypes.c_int, ctypes.c_int]
    _lib.eft_matmul.restype = ctypes.c_int

    # eft_reduce_sum
    _lib.eft_reduce_sum.argtypes = [P_F, P_F, P_F, ctypes.c_int]
    _lib.eft_reduce_sum.restype = ctypes.c_int

    # eft_zero_compensation
    _lib.eft_zero_compensation.argtypes = [P_F, ctypes.c_int]
    _lib.eft_zero_compensation.restype = ctypes.c_int

    # Try to configure compact functions (may not be present in core library)
    _try_configure_compact_functions()


def _try_configure_compact_functions():
    """Try to configure compact adam functions. These may be in a separate library."""
    global _lib

    P_F = ctypes.POINTER(ctypes.c_float)
    P_V = ctypes.c_void_p

    # Check if compact functions already exist in core lib
    if hasattr(_lib, 'compact_adam_create'):
        _configure_compact_signatures()
        return

    # Try to load compact library separately
    compact_candidates = [
        Path(__file__).parent.parent / "bin" / "quarterbit_compact.dll",
        Path(__file__).parent.parent.parent / "build" / "quarterbit_compact.dll",
        Path(r"C:\QuarterBit\build\quarterbit_compact.dll"),
    ]
    if sys.platform != "win32":
        compact_candidates = [
            Path(__file__).parent.parent / "bin" / "libquarterbit_compact.so",
            Path(__file__).parent.parent.parent / "build" / "libquarterbit_compact.so",
        ]

    for path in compact_candidates:
        if path.exists():
            try:
                compact_lib = ctypes.CDLL(str(path))
                # Add compact functions to main lib object for unified access
                _lib.compact_adam_create = compact_lib.compact_adam_create
                _lib.compact_adam_step = compact_lib.compact_adam_step
                _lib.compact_adam_destroy = compact_lib.compact_adam_destroy
                _lib.compact_adam_memory_info = compact_lib.compact_adam_memory_info
                _configure_compact_signatures()
                return
            except Exception:
                pass


def _configure_compact_signatures():
    """Configure ctypes signatures for compact functions."""
    global _lib

    P_F = ctypes.POINTER(ctypes.c_float)
    P_V = ctypes.c_void_p
    P_D = ctypes.POINTER(ctypes.c_double)

    # compact_adam_create(n, use_ef) -> handle
    _lib.compact_adam_create.argtypes = [ctypes.c_int, ctypes.c_int]
    _lib.compact_adam_create.restype = P_V

    # compact_adam_step(handle, params, grads, lr, beta1, beta2, eps, step) -> int
    _lib.compact_adam_step.argtypes = [
        P_V, P_F, P_F,  # handle, params, grads
        ctypes.c_float, ctypes.c_float, ctypes.c_float, ctypes.c_float,  # lr, beta1, beta2, eps
        ctypes.c_int  # step
    ]
    _lib.compact_adam_step.restype = ctypes.c_int

    # compact_adam_destroy(handle)
    _lib.compact_adam_destroy.argtypes = [P_V]
    _lib.compact_adam_destroy.restype = None

    # compact_adam_memory_info(n, use_ef, fp32_bytes, compact_bytes, savings_pct)
    _lib.compact_adam_memory_info.argtypes = [
        ctypes.c_longlong, ctypes.c_int, P_D, P_D, P_D
    ]
    _lib.compact_adam_memory_info.restype = None


# CompactEFT library (separate from compact)
_compact_eft_lib = None

def _load_compact_eft_lib():
    """Load the CompactEFT library (production optimizer with EFT precision)."""
    global _compact_eft_lib

    if _compact_eft_lib is not None:
        return _compact_eft_lib

    P_F = ctypes.POINTER(ctypes.c_float)
    P_V = ctypes.c_void_p

    candidates = [
        Path(__file__).parent.parent / "bin" / "quarterbit_compact_eft.dll",
        Path(__file__).parent.parent.parent / "build" / "quarterbit_compact_eft.dll",
        Path(r"C:\QuarterBit\build\quarterbit_compact_eft.dll"),
    ]
    if sys.platform != "win32":
        candidates = [
            Path(__file__).parent.parent / "bin" / "libquarterbit_compact_eft.so",
            Path(__file__).parent.parent.parent / "build" / "libquarterbit_compact_eft.so",
        ]

    for path in candidates:
        if path.exists():
            try:
                lib = ctypes.CDLL(str(path))

                # Configure signatures
                lib.compact_eft_adam_create.argtypes = [ctypes.c_int]
                lib.compact_eft_adam_create.restype = P_V

                lib.compact_eft_adam_step.argtypes = [
                    P_V, P_F, P_F,
                    ctypes.c_float, ctypes.c_float, ctypes.c_float, ctypes.c_float,
                    ctypes.c_int
                ]
                lib.compact_eft_adam_step.restype = ctypes.c_int

                lib.compact_eft_adam_destroy.argtypes = [P_V]
                lib.compact_eft_adam_destroy.restype = None

                _compact_eft_lib = lib
                return lib
            except Exception as e:
                warnings.warn(f"Failed to load CompactEFT library: {e}")

    return None


def get_lib():
    """Get the loaded library, loading if necessary."""
    if not _load_backend():
        raise RuntimeError(
            "QuarterBit backend not available. "
            "Ensure CUDA is installed and you have a valid license."
        )
    return _lib

def is_available() -> bool:
    """Check if QuarterBit backend is available."""
    return _load_backend()

def get_backend_info() -> dict:
    """Get backend information."""
    if not _load_backend():
        return {"available": False}

    return {
        "available": True,
        "version": "1.0.0",
    }
