//! Python bindings for GPU-accelerated operations.

use std::sync::Arc;

use pyo3::prelude::*;
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use numpy::{PyReadonlyArray2, PyUntypedArrayMethods};

use dyf_core::gpu::{GpuContext, GpuError};
use dyf_core::gpu::tree_routing::GpuTreeRouter;
use dyf_core::gpu::pca::GpuPowerIteration;
use dyf_core::gpu::hash_scatter::GpuHashScatter;

use crate::format_bindings::PyDyfFile;

fn gpu_err(e: GpuError) -> PyErr {
    PyRuntimeError::new_err(e.to_string())
}

/// GPU compute context — wraps wgpu device and queue.
#[pyclass]
pub struct PyGpuContext {
    ctx: Option<GpuContext>,
}

#[pymethods]
impl PyGpuContext {
    /// Initialize GPU context, selecting the best available adapter.
    #[new]
    fn new() -> PyResult<Self> {
        let ctx = GpuContext::new().map_err(gpu_err)?;
        Ok(PyGpuContext { ctx: Some(ctx) })
    }

    /// Human-readable GPU info string.
    fn info(&self) -> PyResult<String> {
        Ok(self.ctx.as_ref().unwrap().info())
    }
}

/// GPU-accelerated tree routing.
#[pyclass]
pub struct PyGpuTreeRouter {
    router: GpuTreeRouter,
}

#[pymethods]
impl PyGpuTreeRouter {
    /// Create a GPU tree router from a DYF file's tree.
    ///
    /// Args:
    ///     ctx: GPU context (consumed — create a new one for other pipelines)
    ///     dyf_file: Open DYF file containing the tree
    #[new]
    fn new(ctx: &mut PyGpuContext, dyf_file: &PyDyfFile) -> PyResult<Self> {
        let gpu_ctx = ctx.ctx.take().ok_or_else(|| {
            PyRuntimeError::new_err("GpuContext already consumed")
        })?;
        let tree = dyf_file.inner.tree();
        let router = GpuTreeRouter::new(Arc::new(gpu_ctx), tree).map_err(gpu_err)?;
        Ok(PyGpuTreeRouter { router })
    }

    /// Route a batch of embeddings through the tree.
    ///
    /// Args:
    ///     embeddings: 2D numpy array, shape (n, dim)
    ///
    /// Returns:
    ///     List of leaf node indices (one per embedding)
    fn route_batch(&self, embeddings: PyReadonlyArray2<'_, f32>) -> PyResult<Vec<u32>> {
        let shape = embeddings.shape();
        let n = shape[0];
        let data = embeddings.as_slice().map_err(|_| {
            PyValueError::new_err("Embeddings array must be contiguous")
        })?;
        self.router.route_batch(data, n).map_err(gpu_err)
    }
}

/// GPU-accelerated power iteration PCA.
///
/// Args:
///     ctx: GPU context (consumed)
///     xtx: Covariance matrix X^T X, 2D numpy array shape (d, d)
///     k: Number of components
///     n_iter: Power iteration steps (default 20)
///     seed: Random seed (default 42)
///
/// Returns:
///     Tuple of (components as list of lists, eigenvalues as list)
#[pyfunction]
#[pyo3(signature = (ctx, xtx, k, n_iter=20, seed=42))]
pub fn py_gpu_power_iteration(
    ctx: &mut PyGpuContext,
    xtx: PyReadonlyArray2<'_, f32>,
    k: usize,
    n_iter: usize,
    seed: u64,
) -> PyResult<(Vec<Vec<f32>>, Vec<f32>)> {
    let gpu_ctx = ctx.ctx.take().ok_or_else(|| {
        PyRuntimeError::new_err("GpuContext already consumed")
    })?;
    let shape = xtx.shape();
    let d = shape[0];
    if shape[1] != d {
        return Err(PyValueError::new_err("xtx must be square"));
    }
    let data = xtx.as_slice().map_err(|_| {
        PyValueError::new_err("xtx must be contiguous")
    })?;

    let pca = GpuPowerIteration::new(Arc::new(gpu_ctx)).map_err(gpu_err)?;
    let (flat_components, eigenvalues) = pca
        .power_iteration(data, d, k, n_iter, seed)
        .map_err(gpu_err)?;

    // Reshape k*d flat → list of k vectors of d
    let components: Vec<Vec<f32>> = flat_components
        .chunks(d)
        .map(|c| c.to_vec())
        .collect();

    Ok((components, eigenvalues))
}

/// GPU-accelerated LSH hash computation.
///
/// Args:
///     ctx: GPU context (consumed)
///     embeddings: 2D numpy array, shape (n, dim)
///     hyperplanes: 2D numpy array, shape (num_bits, dim)
///
/// Returns:
///     List of u32 hash values
#[pyfunction]
pub fn py_gpu_hash_batch(
    ctx: &mut PyGpuContext,
    embeddings: PyReadonlyArray2<'_, f32>,
    hyperplanes: PyReadonlyArray2<'_, f32>,
) -> PyResult<Vec<u32>> {
    let gpu_ctx = ctx.ctx.take().ok_or_else(|| {
        PyRuntimeError::new_err("GpuContext already consumed")
    })?;

    let emb_shape = embeddings.shape();
    let hp_shape = hyperplanes.shape();
    let n = emb_shape[0];
    let dim = emb_shape[1];
    let num_bits = hp_shape[0];

    if hp_shape[1] != dim {
        return Err(PyValueError::new_err(
            "hyperplanes dim must match embeddings dim",
        ));
    }
    if num_bits > 32 {
        return Err(PyValueError::new_err("GPU supports up to 32 hash bits"));
    }

    let emb_data = embeddings.as_slice().map_err(|_| {
        PyValueError::new_err("Embeddings must be contiguous")
    })?;
    let hp_data = hyperplanes.as_slice().map_err(|_| {
        PyValueError::new_err("Hyperplanes must be contiguous")
    })?;

    let hasher = GpuHashScatter::new(Arc::new(gpu_ctx)).map_err(gpu_err)?;
    hasher.hash_batch(emb_data, hp_data, n, dim, num_bits).map_err(gpu_err)
}
