use crate::error::{TuphiumError, TuphiumResult};
use cublas;
use cublas::api::Operation;
use serde::de::{self, Deserialize, Deserializer, Visitor};
use serde::ser::{Serialize, SerializeStruct, Serializer};
use rustacuda::memory::DeviceBuffer;
use rustacuda::prelude::*;
use std::fmt;
use std::sync::OnceLock;

// ── Global CUDA + cuBLAS initialisation ──

// Wrap cuBLAS context in Send+Sync wrapper (Context is Sync but not Send)
struct CublasCtx(cublas::Context);
unsafe impl Send for CublasCtx {}
unsafe impl Sync for CublasCtx {}

fn try_cublas_ctx() -> TuphiumResult<&'static cublas::Context> {
    static CUBLAS: OnceLock<CublasCtx> = OnceLock::new();
    if let Some(ctx) = CUBLAS.get() {
        return Ok(&ctx.0);
    }

    rustacuda::init(CudaFlags::empty()).map_err(|error| TuphiumError::BackendUnavailable {
        operation: "CUDA::init",
        backend: "CUDA".to_string(),
        message: format!("{error:?}"),
    })?;
    let device = Device::get_device(0).map_err(|error| TuphiumError::BackendUnavailable {
        operation: "CUDA::get_device",
        backend: "CUDA".to_string(),
        message: format!("{error:?}"),
    })?;
    let ctx = Context::create_and_push(ContextFlags::MAP_HOST | ContextFlags::SCHED_AUTO, device)
        .map_err(|error| TuphiumError::BackendUnavailable {
        operation: "CUDA::create_context",
        backend: "CUDA".to_string(),
        message: format!("{error:?}"),
    })?;
    Box::leak(Box::new(ctx));
    let cublas_context =
        cublas::Context::new().map_err(|error| TuphiumError::BackendUnavailable {
            operation: "cuBLAS::new",
            backend: "cuBLAS".to_string(),
            message: format!("{error:?}"),
        })?;

    if CUBLAS.set(CublasCtx(cublas_context)).is_err() {
        return CUBLAS
            .get()
            .map(|ctx| &ctx.0)
            .ok_or_else(|| TuphiumError::BackendUnavailable {
                operation: "cuBLAS::new",
                backend: "cuBLAS".to_string(),
                message: "context initialization race failed".to_string(),
            });
    }

    Ok(&CUBLAS.get().expect("cuBLAS context was just initialized").0)
}

// ── GPU-accelerated Matrix ──
#[derive(Debug)]
pub struct Matrix {
    pub rows: usize,
    pub cols: usize,
    pub data: Vec<f32>,
    device: Option<DeviceBuffer<f32>>,
    device_dirty: bool,
}

impl Clone for Matrix {
    fn clone(&self) -> Self {
        Self {
            rows: self.rows,
            cols: self.cols,
            data: self.data.clone(),
            device: None,
            device_dirty: true,
        }
    }
}

impl Default for Matrix {
    fn default() -> Self {
        Self {
            rows: 0,
            cols: 0,
            data: vec![],
            device: None,
            device_dirty: false,
        }
    }
}

impl Serialize for Matrix {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let mut s = serializer.serialize_struct("Matrix", 3)?;
        s.serialize_field("rows", &self.rows)?;
        s.serialize_field("cols", &self.cols)?;
        s.serialize_field("data", &self.data)?;
        s.end()
    }
}

impl<'de> Deserialize<'de> for Matrix {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        struct MatrixVisitor;

        impl<'de> Visitor<'de> for MatrixVisitor {
            type Value = Matrix;

            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
                f.write_str("struct Matrix")
            }

            fn visit_map<A: de::MapAccess<'de>>(self, mut map: A) -> Result<Matrix, A::Error> {
                let mut rows = None;
                let mut cols = None;
                let mut data = None;

                while let Some(key) = map.next_key::<String>()? {
                    match key.as_str() {
                        "rows" => rows = Some(map.next_value()?),
                        "cols" => cols = Some(map.next_value()?),
                        "data" => data = Some(map.next_value()?),
                        _ => {
                            let _: de::IgnoredAny = map.next_value()?;
                        }
                    }
                }

                let rows = rows.ok_or_else(|| de::Error::missing_field("rows"))?;
                let cols = cols.ok_or_else(|| de::Error::missing_field("cols"))?;
                let data = data.ok_or_else(|| de::Error::missing_field("data"))?;

                Ok(Matrix {
                    rows,
                    cols,
                    data,
                    device: None,
                    device_dirty: true,
                })
            }
        }

        deserializer.deserialize_struct("Matrix", &["rows", "cols", "data"], MatrixVisitor)
    }
}

impl Matrix {
    /// Creates a `rows × cols` matrix filled with zeros.
    ///
    /// The GPU buffer is allocated lazily on the first operation that needs it.
    pub fn zeros(rows: usize, cols: usize) -> TuphiumResult<Self> {
        let len = rows
            .checked_mul(cols)
            .ok_or(TuphiumError::DimensionOverflow {
                operation: "Matrix::zeros",
                rows,
                cols,
            })?;
        Ok(Self {
            rows,
            cols,
            data: vec![0.0; len],
            device: None,
            device_dirty: true,
        })
    }

    /// Builds a `rows × cols` matrix from a flat, row-major `data` buffer.
    ///
    /// Element `(r, c)` is read from `data[r * cols + c]`. The device copy is
    /// uploaded lazily (and transposed to column-major for cuBLAS) on first use.
    ///
    /// # Panics
    ///
    /// Panics if `data.len()` is not exactly `rows * cols`.
    pub fn from_data(rows: usize, cols: usize, data: Vec<f32>) -> TuphiumResult<Self> {
        let expected = rows
            .checked_mul(cols)
            .ok_or(TuphiumError::DimensionOverflow {
                operation: "Matrix::from_data",
                rows,
                cols,
            })?;
        if data.len() != expected {
            return Err(TuphiumError::LengthMismatch {
                operation: "Matrix::from_data",
                expected,
                actual: data.len(),
            });
        }
        Ok(Self {
            rows,
            cols,
            data,
            device: None,
            device_dirty: true,
        })
    }

    /// Returns the host-side element at row `r`, column `c`.
    ///
    /// Reads from `data` only; call [`Matrix::sync_from_device`] first if the
    /// latest values may still live on the GPU.
    ///
    /// # Panics
    ///
    /// Panics if the index falls outside the backing buffer.
    #[inline]
    pub fn get(&self, r: usize, c: usize) -> TuphiumResult<f32> {
        let index = self.checked_index(r, c, "Matrix::get")?;
        Ok(self.data[index])
    }

    /// Writes `v` into the host-side row `r`, column `c`.
    ///
    /// Does not mark the device buffer dirty; call [`Matrix::mark_dirty`] if the
    /// change must reach the GPU.
    ///
    /// # Panics
    ///
    /// Panics if the index falls outside the backing buffer.
    #[inline]
    pub fn set(&mut self, r: usize, c: usize, v: f32) -> TuphiumResult<()> {
        let index = self.checked_index(r, c, "Matrix::set")?;
        self.data[index] = v;
        Ok(())
    }

    /// Returns row `i` as a contiguous host-side slice of length `cols`.
    pub fn row(&self, i: usize) -> TuphiumResult<&[f32]> {
        if i >= self.rows {
            return Err(TuphiumError::RowIndexOutOfBounds {
                operation: "Matrix::row",
                row: i,
                rows: self.rows,
            });
        }
        let start = i * self.cols;
        Ok(&self.data[start..start + self.cols])
    }

    // host <-> column-major helpers
    fn to_col_major(&self) -> Vec<f32> {
        let n = self.rows * self.cols;
        let mut cm = vec![0.0; n];
        for r in 0..self.rows {
            for c in 0..self.cols {
                cm[c * self.rows + r] = self.data[r * self.cols + c];
            }
        }
        cm
    }

    fn from_col_major(&mut self, cm: &[f32]) {
        for r in 0..self.rows {
            for c in 0..self.cols {
                self.data[r * self.cols + c] = cm[c * self.rows + r];
            }
        }
    }

    fn try_ensure_device(&mut self) -> TuphiumResult<&mut DeviceBuffer<f32>> {
        try_cublas_ctx()?;
        if self.device_dirty || self.device.is_none() {
            let cm = self.to_col_major();
            match self.device {
                Some(ref mut buf) => {
                    buf.copy_from(&cm)
                        .map_err(|error| TuphiumError::BackendOperationFailed {
                            operation: "Matrix::sync_to_device",
                            backend: "CUDA".to_string(),
                            message: format!("{error:?}"),
                        })?;
                }
                None => {
                    let buf = DeviceBuffer::from_slice(&cm).map_err(|error| {
                        TuphiumError::BackendOperationFailed {
                            operation: "Matrix::sync_to_device",
                            backend: "CUDA".to_string(),
                            message: format!("{error:?}"),
                        }
                    })?;
                    self.device = Some(buf);
                }
            }
            self.device_dirty = false;
        }
        self.device
            .as_mut()
            .ok_or_else(|| TuphiumError::InvalidState {
                operation: "Matrix::sync_to_device",
                state: "device buffer was not initialized",
            })
    }

    /// Force a host → device copy, even if the device buffer is not dirty.
    pub fn sync_to_device(&mut self) -> TuphiumResult<()> {
        self.device_dirty = true;
        self.try_ensure_device()?;
        Ok(())
    }

    /// Force a device → host copy, bringing GPU-resident data back to `self.data`.
    pub fn sync_from_device(&mut self) -> TuphiumResult<()> {
        if let Some(ref buf) = self.device {
            let mut cm = vec![0.0; self.rows * self.cols];
            buf.copy_to(&mut cm)
                .map_err(|error| TuphiumError::BackendOperationFailed {
                    operation: "Matrix::sync_from_device",
                    backend: "CUDA".to_string(),
                    message: format!("{error:?}"),
                })?;
            self.from_col_major(&cm);
        }
        self.device_dirty = false;
        Ok(())
    }

    /// Mark the host-side `data` as modified, so the next GPU operation re-uploads it.
    pub fn mark_dirty(&mut self) {
        self.device_dirty = true;
    }

    // ── cuBLAS operations (mapped via gemm) ──

    /// Matrix–vector product `y = self · x`, computed on the GPU via cuBLAS gemm.
    ///
    /// Takes `&mut self` because the matrix is uploaded to the device on demand.
    /// `x` must have length `cols`; the result has length `rows`.
    ///
    /// # Panics
    ///
    /// Panics if `x.len() != cols`.
    pub fn matvec(&mut self, x: &[f32]) -> TuphiumResult<Vec<f32>> {
        if x.len() != self.cols {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Matrix::matvec",
                expected: vec![self.cols],
                actual: vec![x.len()],
            });
        }
        let rows = self.rows;
        let cols = self.cols;
        if rows == 0 || cols == 0 {
            return Ok(vec![0.0; rows]);
        }

        let ctx = try_cublas_ctx()?;
        let d_a = self.try_ensure_device()?;
        let mut d_x =
            DeviceBuffer::from_slice(x).map_err(|error| TuphiumError::BackendOperationFailed {
                operation: "Matrix::matvec",
                backend: "CUDA".to_string(),
                message: format!("{error:?}"),
            })?;
        let mut d_y = DeviceBuffer::from_slice(&vec![0.0; rows]).map_err(|error| {
            TuphiumError::BackendOperationFailed {
                operation: "Matrix::matvec",
                backend: "CUDA".to_string(),
                message: format!("{error:?}"),
            }
        })?;

        let m = rows as i32;
        let n = 1i32;
        let k = cols as i32;

        let mut alpha: f32 = 1.0;
        let mut beta: f32 = 0.0;

        ctx.gemm(
            Operation::NoTrans,
            Operation::NoTrans,
            m,
            n,
            k,
            &mut alpha,
            d_a.as_mut_ptr(),
            m,
            d_x.as_mut_ptr(),
            k,
            &mut beta,
            d_y.as_mut_ptr(),
            m,
        )
        .map_err(|error| TuphiumError::BackendOperationFailed {
            operation: "Matrix::matvec",
            backend: "cuBLAS".to_string(),
            message: format!("{error:?}"),
        })?;

        let mut result = vec![0.0; rows];
        d_y.copy_to(&mut result)
            .map_err(|error| TuphiumError::BackendOperationFailed {
                operation: "Matrix::matvec",
                backend: "CUDA".to_string(),
                message: format!("{error:?}"),
            })?;
        Ok(result)
    }

    /// Transposed matrix–vector product `y = selfᵀ · v`, via cuBLAS gemm.
    ///
    /// Takes `&mut self` because the matrix is uploaded to the device on demand.
    /// `v` must have length `rows`; the result has length `cols`. Used in the
    /// backward pass to propagate gradients to a layer's inputs.
    ///
    /// # Panics
    ///
    /// Panics if `v.len() != rows`.
    pub fn t_matvec(&mut self, v: &[f32]) -> TuphiumResult<Vec<f32>> {
        if v.len() != self.rows {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Matrix::t_matvec",
                expected: vec![self.rows],
                actual: vec![v.len()],
            });
        }
        let rows = self.rows;
        let cols = self.cols;
        if rows == 0 || cols == 0 {
            return Ok(vec![0.0; cols]);
        }

        let ctx = try_cublas_ctx()?;
        let d_a = self.try_ensure_device()?;
        let mut d_v =
            DeviceBuffer::from_slice(v).map_err(|error| TuphiumError::BackendOperationFailed {
                operation: "Matrix::t_matvec",
                backend: "CUDA".to_string(),
                message: format!("{error:?}"),
            })?;
        let mut d_y = DeviceBuffer::from_slice(&vec![0.0; cols]).map_err(|error| {
            TuphiumError::BackendOperationFailed {
                operation: "Matrix::t_matvec",
                backend: "CUDA".to_string(),
                message: format!("{error:?}"),
            }
        })?;

        let m = cols as i32;
        let n = 1i32;
        let k = rows as i32;

        let mut alpha: f32 = 1.0;
        let mut beta: f32 = 0.0;

        ctx.gemm(
            Operation::Trans,
            Operation::NoTrans,
            m,
            n,
            k,
            &mut alpha,
            d_a.as_mut_ptr(),
            k,
            d_v.as_mut_ptr(),
            k,
            &mut beta,
            d_y.as_mut_ptr(),
            m,
        )
        .map_err(|error| TuphiumError::BackendOperationFailed {
            operation: "Matrix::t_matvec",
            backend: "cuBLAS".to_string(),
            message: format!("{error:?}"),
        })?;

        let mut result = vec![0.0; cols];
        d_y.copy_to(&mut result)
            .map_err(|error| TuphiumError::BackendOperationFailed {
                operation: "Matrix::t_matvec",
                backend: "CUDA".to_string(),
                message: format!("{error:?}"),
            })?;
        Ok(result)
    }

    /// Accumulates the outer product into the matrix: `self += a ⊗ b`, via gemm.
    ///
    /// `a` indexes rows (length `rows`) and `b` indexes columns (length `cols`).
    /// The result accumulates directly in the device buffer (cuBLAS `beta = 1`);
    /// call [`Matrix::sync_from_device`] before reading `data`. Used in the
    /// backward pass to accumulate weight gradients (`dz ⊗ x`).
    ///
    /// # Panics
    ///
    /// Panics if `a.len() != rows` or `b.len() != cols`.
    pub fn add_outer(&mut self, a: &[f32], b: &[f32]) -> TuphiumResult<()> {
        if a.len() != self.rows || b.len() != self.cols {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Matrix::add_outer",
                expected: vec![self.rows, self.cols],
                actual: vec![a.len(), b.len()],
            });
        }
        if self.rows == 0 || self.cols == 0 {
            return Ok(());
        }

        let m = self.rows as i32;
        let n = self.cols as i32;
        let ctx = try_cublas_ctx()?;
        let mut d_a_vec =
            DeviceBuffer::from_slice(a).map_err(|error| TuphiumError::BackendOperationFailed {
                operation: "Matrix::add_outer",
                backend: "CUDA".to_string(),
                message: format!("{error:?}"),
            })?;
        let mut d_b_vec =
            DeviceBuffer::from_slice(b).map_err(|error| TuphiumError::BackendOperationFailed {
                operation: "Matrix::add_outer",
                backend: "CUDA".to_string(),
                message: format!("{error:?}"),
            })?;
        let d_mat = self.try_ensure_device()?;

        let k = 1i32;

        let mut alpha: f32 = 1.0;
        let mut beta: f32 = 1.0;

        ctx.gemm(
            Operation::NoTrans,
            Operation::Trans,
            m,
            n,
            k,
            &mut alpha,
            d_a_vec.as_mut_ptr(),
            m,
            d_b_vec.as_mut_ptr(),
            n,
            &mut beta,
            d_mat.as_mut_ptr(),
            m,
        )
        .map_err(|error| TuphiumError::BackendOperationFailed {
            operation: "Matrix::add_outer",
            backend: "cuBLAS".to_string(),
            message: format!("{error:?}"),
        })?;
        Ok(())
    }

    /// Overwrites every host element with `val` and invalidates the device copy
    /// so the next GPU operation re-uploads it.
    pub fn fill(&mut self, val: f32) {
        for x in self.data.iter_mut() {
            *x = val;
        }
        self.device = None;
        self.device_dirty = true;
    }

    fn checked_index(
        &self,
        row: usize,
        col: usize,
        operation: &'static str,
    ) -> TuphiumResult<usize> {
        if row >= self.rows || col >= self.cols {
            return Err(TuphiumError::IndexOutOfBounds {
                operation,
                row,
                col,
                rows: self.rows,
                cols: self.cols,
            });
        }
        row.checked_mul(self.cols)
            .and_then(|base| base.checked_add(col))
            .filter(|&index| index < self.data.len())
            .ok_or(TuphiumError::IndexOutOfBounds {
                operation,
                row,
                col,
                rows: self.rows,
                cols: self.cols,
            })
    }
}
