//! CUDA テンソル

use crate::autograd::GradFn;
use crate::buffer_pool::{pool_acquire, pool_release, pool_print_stats};
use crate::cuda_sys::{self, cudaMemcpyKind, CUDA_SUCCESS};
use crate::device::{get_device, CudaDevice};
use crate::{shape_to_bytes, DType};
use std::cell::UnsafeCell;
use std::ffi::c_void;
use std::sync::Arc;
use tl_backend::BackendResult;

/// Arc ベースのテンソル参照（V5.0 メモリ管理）
/// GradFn 内で入力テンソルの生存を保証する。
pub type TensorRef = Arc<UnsafeCell<CudaTensor>>;

/// 生ポインタから TensorRef を取得する（RC+1）。
pub unsafe fn tensor_ref_from_ptr(ptr: *mut CudaTensor) -> TensorRef {
    let arc: TensorRef = Arc::from_raw(ptr as *const UnsafeCell<CudaTensor>);
    let cloned = arc.clone(); // RC+1
    std::mem::forget(arc); // 元のポインタは生かす
    cloned
}

/// TensorRef から内部の CudaTensor への不変参照を取得する。
#[inline]
pub unsafe fn tensor_ref_get(r: &TensorRef) -> &CudaTensor {
    &*r.get()
}

/// TensorRef から内部の CudaTensor への可変参照を取得する。
#[inline]
pub unsafe fn tensor_ref_get_mut(r: &TensorRef) -> &mut CudaTensor {
    &mut *r.get()
}

/// CUDA GPU バッファ
/// Drop 時に cudaFree を呼んで GPU メモリを解放する。
pub struct CudaBuffer {
    ptr: *mut c_void,
    size: usize,
}

unsafe impl Send for CudaBuffer {}
unsafe impl Sync for CudaBuffer {}

impl CudaBuffer {
    /// 新しい GPU バッファを確保
    pub fn new(size: usize) -> Result<Self, String> {
        let device = get_device();
        let ptr = device.allocate_buffer(size)?;
        Ok(CudaBuffer { ptr, size })
    }

    /// GPU メモリへのポインタ
    pub fn ptr(&self) -> *mut c_void {
        self.ptr
    }

    /// バッファサイズ（バイト数）
    pub fn size(&self) -> usize {
        self.size
    }
}

impl Drop for CudaBuffer {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            unsafe {
                cuda_sys::cudaFree(self.ptr);
            }
        }
    }
}

/// Autograd メタデータ
pub struct AutogradMeta {
    /// 勾配（累積）
    pub grad: Option<CudaTensor>,
    /// 勾配関数（リーフノードは None）
    pub grad_fn: Option<Box<dyn GradFn>>,
    /// 勾配が必要か
    pub requires_grad: bool,
}

/// CUDA GPU テンソル
pub struct CudaTensor {
    /// GPU バッファ
    pub(crate) buffer: Arc<CudaBuffer>,
    /// 形状
    pub shape: Vec<usize>,
    /// データ型
    pub dtype: DType,
    /// デバイス
    #[allow(dead_code)]
    device: Arc<CudaDevice>,
    /// Autograd メタデータ（None = autograd 不要）
    pub autograd: Option<Box<AutogradMeta>>,
}

impl std::fmt::Debug for CudaTensor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaTensor")
            .field("shape", &self.shape)
            .field("dtype", &self.dtype)
            .field("buffer_ptr", &self.buffer.ptr())
            .finish()
    }
}

impl Clone for CudaTensor {
    fn clone(&self) -> Self {
        // データを複製して新しいテンソルを作成
        self.clone_data().expect("Clone failed")
    }
}

unsafe impl Send for CudaTensor {}
unsafe impl Sync for CudaTensor {}

impl CudaTensor {
    /// 新しいテンソルを作成（未初期化）
    pub fn uninit(shape: &[usize], dtype: DType) -> Self {
        let device = get_device();
        let size = shape_to_bytes(shape, dtype);

        // プールから取得を試みる
        let buffer = pool_acquire(size).unwrap_or_else(|| {
            // プールになければ新規確保
            match CudaBuffer::new(size) {
                Ok(b) => Arc::new(b),
                Err(e) => {
                    eprintln!("CUDA OOM at size {}. Printing Pool Stats:", size);
                    pool_print_stats();
                    let make = crate::ffi_ops::FFI_COUNTERS.make.load(std::sync::atomic::Ordering::Relaxed);
                    let release = crate::ffi_ops::FFI_COUNTERS.release.load(std::sync::atomic::Ordering::Relaxed);
                    let acquire = crate::ffi_ops::FFI_COUNTERS.acquire.load(std::sync::atomic::Ordering::Relaxed);
                    eprintln!("FFI_COUNTERS: make={}, release={}, acquire={}, live={}", make, release, acquire, make + acquire - release);
                    panic!("CUDA buffer allocation failed: {}", e);
                }
            }
        });

        CudaTensor {
            buffer,
            shape: shape.to_vec(),
            dtype,
            device,
            autograd: None,
        }
    }

    /// GPU buffer を共有して shape だけ変更した新テンソルを作成（ゼロコピー）
    pub(crate) fn view_with_shape(&self, new_shape: &[usize]) -> Self {
        CudaTensor {
            buffer: self.buffer.clone(), // Arc clone = 参照カウント+1
            shape: new_shape.to_vec(),
            dtype: self.dtype,
            device: self.device.clone(),
            autograd: None,
        }
    }

    /// ゼロで初期化されたテンソルを作成
    pub fn zeros(shape: &[usize], dtype: DType) -> Self {
        let tensor = Self::uninit(shape, dtype);
        let size = shape_to_bytes(shape, dtype);
        if size > 0 {
            let err = unsafe { cuda_sys::cudaMemset(tensor.buffer.ptr(), 0, size) };
            if err != CUDA_SUCCESS {
                eprintln!("cudaMemset failed in zeros(): {}", err);
            }
        }
        tensor
    }

    /// スライスからテンソルを作成（Host→Device コピー）
    pub fn from_slice<T: Copy>(data: &[T], shape: &[usize], dtype: DType) -> Self {
        let tensor = Self::uninit(shape, dtype);
        let byte_size = data.len() * std::mem::size_of::<T>();
        if byte_size > 0 {
            let err = unsafe {
                cuda_sys::cudaMemcpy(
                    tensor.buffer.ptr(),
                    data.as_ptr() as *const c_void,
                    byte_size,
                    cudaMemcpyKind::cudaMemcpyHostToDevice,
                )
            };
            if err != CUDA_SUCCESS {
                panic!("cudaMemcpy HostToDevice failed in from_slice(): {}", err);
            }
        }
        tensor
    }

    /// 形状
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    /// データ型
    pub fn dtype(&self) -> DType {
        self.dtype
    }

    /// 要素数
    pub fn elem_count(&self) -> usize {
        self.shape.iter().product()
    }

    /// GPU バッファへのポインタ
    pub fn buffer_ptr(&self) -> *mut c_void {
        self.buffer.ptr()
    }

    /// GPU バッファの Arc を取得
    pub fn buffer_arc(&self) -> &Arc<CudaBuffer> {
        &self.buffer
    }

    /// 全て1で初期化
    pub fn ones(shape: &[usize], dtype: DType) -> Self {
        let count = shape.iter().product::<usize>();
        match dtype {
            DType::F32 => {
                let data = vec![1.0f32; count];
                Self::from_slice(&data, shape, dtype)
            }
            DType::I64 => {
                let data = vec![1i64; count];
                Self::from_slice(&data, shape, dtype)
            }
            DType::I32 => {
                let data = vec![1i32; count];
                Self::from_slice(&data, shape, dtype)
            }
            DType::F64 => {
                let data = vec![1.0f64; count];
                Self::from_slice(&data, shape, dtype)
            }
            DType::U8 => {
                let data = vec![1u8; count];
                Self::from_slice(&data, shape, dtype)
            }
            DType::U32 => {
                let data = vec![1u32; count];
                Self::from_slice(&data, shape, dtype)
            }
            DType::F16 => {
                let data = vec![half::f16::ONE; count];
                Self::from_slice(&data, shape, dtype)
            }
            DType::BF16 => {
                let data = vec![half::bf16::ONE; count];
                Self::from_slice(&data, shape, dtype)
            }
        }
    }

    /// 正規乱数で初期化
    pub fn randn(shape: &[usize], dtype: DType) -> Self {
        use rand::Rng;
        let count = shape.iter().product::<usize>();
        match dtype {
            DType::F32 => {
                let mut rng = rand::thread_rng();
                let data: Vec<f32> = (0..count)
                    .map(|_| {
                        let u1: f32 = 1.0 - rng.gen::<f32>();
                        let u2: f32 = rng.gen();
                        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos()
                    })
                    .collect();
                Self::from_slice(&data, shape, dtype)
            }
            DType::F64 => {
                let mut rng = rand::thread_rng();
                let data: Vec<f64> = (0..count)
                    .map(|_| {
                        let u1: f64 = 1.0 - rng.gen::<f64>();
                        let u2: f64 = rng.gen();
                        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
                    })
                    .collect();
                Self::from_slice(&data, shape, dtype)
            }
            DType::F16 => {
                let mut rng = rand::thread_rng();
                let data: Vec<half::f16> = (0..count)
                    .map(|_| {
                        let u1: f32 = 1.0 - rng.gen::<f32>();
                        let u2: f32 = rng.gen();
                        let f = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
                        half::f16::from_f32(f)
                    })
                    .collect();
                Self::from_slice(&data, shape, dtype)
            }
            DType::BF16 => {
                let mut rng = rand::thread_rng();
                let data: Vec<half::bf16> = (0..count)
                    .map(|_| {
                        let u1: f32 = 1.0 - rng.gen::<f32>();
                        let u2: f32 = rng.gen();
                        let f = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
                        half::bf16::from_f32(f)
                    })
                    .collect();
                Self::from_slice(&data, shape, dtype)
            }
            DType::U8 | DType::U32 | DType::I32 | DType::I64 => {
                let mut rng = rand::thread_rng();
                eprintln!("Warning: randn for {:?} requested, mathematical normal distribution for integers is not well-defined. Returning zeros.", dtype);
                Self::zeros(shape, dtype)
            }
        }
    }

    /// データを CPU にコピー
    pub fn to_vec<T: Copy + Default>(&self) -> Vec<T> {
        // GPU 演算の完了を保証
        crate::stream::sync_stream();

        let count = self.elem_count();
        if count == 0 {
            return Vec::new();
        }

        let t_size = std::mem::size_of::<T>();
        let dtype_size = self.dtype.size_in_bytes();

        if t_size == dtype_size {
            // 直接コピー（型サイズが一致）
            let byte_size = count * t_size;
            let mut result = vec![T::default(); count];
            let err = unsafe {
                cuda_sys::cudaMemcpy(
                    result.as_mut_ptr() as *mut c_void,
                    self.buffer.ptr() as *const c_void,
                    byte_size,
                    cudaMemcpyKind::cudaMemcpyDeviceToHost,
                )
            };
            if err != CUDA_SUCCESS {
                eprintln!("cudaMemcpy DeviceToHost failed in to_vec(): {} (byte_size={}, buffer_size={}, shape={:?}, dtype={:?})",
                    err, byte_size, self.buffer.size(), self.shape, self.dtype);
                return vec![T::default(); count];
            }
            result
        } else if dtype_size == 4 && t_size == 8 {
            // F32/I32 → i64: まず f32 で読んでから変換
            let byte_size = count * 4;
            let mut f32_buf = vec![0.0f32; count];
            let err = unsafe {
                cuda_sys::cudaMemcpy(
                    f32_buf.as_mut_ptr() as *mut c_void,
                    self.buffer.ptr() as *const c_void,
                    byte_size,
                    cudaMemcpyKind::cudaMemcpyDeviceToHost,
                )
            };
            if err != CUDA_SUCCESS {
                eprintln!("cudaMemcpy DeviceToHost failed in to_vec() (f32→i64): {} (byte_size={}, buffer_size={}, shape={:?})",
                    err, byte_size, self.buffer.size(), self.shape);
                return vec![T::default(); count];
            }
            // f32 → i64 変換
            let i64_buf: Vec<i64> = f32_buf.iter().map(|&x| x as i64).collect();
            let mut result = vec![T::default(); count];
            unsafe {
                std::ptr::copy_nonoverlapping(
                    i64_buf.as_ptr() as *const T,
                    result.as_mut_ptr(),
                    count,
                );
            }
            result
        } else if dtype_size == 8 && t_size == 4 {
            // I64 → f32: まず i64 で読んでから変換
            let byte_size = count * 8;
            let mut i64_buf = vec![0i64; count];
            let err = unsafe {
                cuda_sys::cudaMemcpy(
                    i64_buf.as_mut_ptr() as *mut c_void,
                    self.buffer.ptr() as *const c_void,
                    byte_size,
                    cudaMemcpyKind::cudaMemcpyDeviceToHost,
                )
            };
            if err != CUDA_SUCCESS {
                eprintln!("cudaMemcpy DeviceToHost failed in to_vec() (i64→f32): {} (byte_size={}, buffer_size={}, shape={:?})",
                    err, byte_size, self.buffer.size(), self.shape);
                return vec![T::default(); count];
            }
            // i64 → f32 変換
            let f32_buf: Vec<f32> = i64_buf.iter().map(|&x| x as f32).collect();
            let mut result = vec![T::default(); count];
            unsafe {
                std::ptr::copy_nonoverlapping(
                    f32_buf.as_ptr() as *const T,
                    result.as_mut_ptr(),
                    count,
                );
            }
            result
        } else {
            // フォールバック: バッファサイズ分だけ読む
            let copy_size = self.buffer.size().min(count * t_size);
            let mut result = vec![T::default(); count];
            let err = unsafe {
                cuda_sys::cudaMemcpy(
                    result.as_mut_ptr() as *mut c_void,
                    self.buffer.ptr() as *const c_void,
                    copy_size,
                    cudaMemcpyKind::cudaMemcpyDeviceToHost,
                )
            };
            if err != CUDA_SUCCESS {
                eprintln!(
                    "cudaMemcpy DeviceToHost failed in to_vec() (fallback): {}",
                    err
                );
            }
            result
        }
    }

    /// データを GPU 上で完全にコピー（DeviceToDevice）
    pub fn clone_data(&self) -> BackendResult<CudaTensor> {
        crate::stream::sync_stream();

        let result = CudaTensor::uninit(self.shape(), self.dtype());
        let size = shape_to_bytes(self.shape(), self.dtype());
        if size > 0 {
            let err = unsafe {
                cuda_sys::cudaMemcpy(
                    result.buffer.ptr(),
                    self.buffer.ptr() as *const c_void,
                    size,
                    cudaMemcpyKind::cudaMemcpyDeviceToDevice,
                )
            };
            if err != CUDA_SUCCESS {
                return Err(tl_backend::BackendError::InternalError(format!(
                    "cudaMemcpy DeviceToDevice failed: {}",
                    err
                )));
            }
        }
        Ok(result)
    }

    /// GPU バッファを共有する浅いクローン（データコピーなし）
    /// autograd メタデータはコピーしない
    pub fn shallow_clone(&self) -> Self {
        CudaTensor {
            buffer: Arc::clone(&self.buffer),
            shape: self.shape.clone(),
            dtype: self.dtype,
            device: Arc::clone(&self.device),
            autograd: if self.requires_grad() {
                Some(Box::new(AutogradMeta {
                    grad: None,
                    grad_fn: None,
                    requires_grad: true,
                }))
            } else {
                None
            },
        }
    }

    /// 既存バッファから新しい形状でテンソルを作成（バッファ共有）
    pub fn from_buffer_shared(shape: Vec<usize>, dtype: DType) -> Self {
        let device = get_device();
        let size = shape_to_bytes(&shape, dtype);
        let buffer = match CudaBuffer::new(size) {
            Ok(b) => Arc::new(b),
            Err(e) => {
                eprintln!("CUDA OOM at size {}. Printing Pool Stats:", size);
                pool_print_stats();
                panic!("CUDA buffer allocation failed: {}", e);
            }
        };
        CudaTensor {
            buffer,
            shape,
            dtype,
            device,
            autograd: None,
        }
    }

    // ========== Autograd メソッド ==========

    /// requires_grad フラグを確認
    pub fn requires_grad(&self) -> bool {
        self.autograd.as_ref().map_or(false, |a| a.requires_grad)
    }

    /// requires_grad を有効化
    pub fn enable_grad(&mut self) {
        if self.autograd.is_none() {
            self.autograd = Some(Box::new(AutogradMeta {
                grad: None,
                grad_fn: None,
                requires_grad: true,
            }));
        } else {
            self.autograd.as_mut().unwrap().requires_grad = true;
        }
    }

    /// grad_fn をセット
    pub fn set_grad_fn(&mut self, grad_fn: Box<dyn GradFn>) {
        self.autograd = Some(Box::new(AutogradMeta {
            grad: None,
            grad_fn: Some(grad_fn),
            requires_grad: true,
        }));
    }

    /// 勾配を取得（shallow clone）
    pub fn get_grad(&self) -> Option<CudaTensor> {
        self.autograd
            .as_ref()
            .and_then(|a| a.grad.as_ref().map(|g| g.shallow_clone()))
    }

    /// 勾配をゼロクリア
    pub fn zero_grad(&mut self) {
        if let Some(ref mut a) = self.autograd {
            a.grad = None;
        }
    }

    /// 勾配を累積
    pub fn accumulate_grad(&mut self, grad: CudaTensor) -> BackendResult<()> {
        if let Some(ref mut meta) = self.autograd {
            let detached_grad = grad.detach();
            if let Some(ref mut g) = meta.grad {
                *g = g.add_impl(&detached_grad)?;
            } else {
                meta.grad = Some(detached_grad);
            }
        }
        Ok(())
    }

    /// backward（逆伝播）— Metal 版と同一構造
    /// 1. DFS でトポロジカル順序を構築
    /// 2. 出力テンソルの勾配を ones で初期化
    /// 3. 逆トポロジカル順序で勾配を伝播
    /// 4. 全ノードの grad_fn をクリア
    pub fn backward(&mut self) -> BackendResult<()> {
        if !self.requires_grad() {
            return Ok(());
        }
        let self_ptr = self as *mut Self;

        // 1. DFS でトポロジカル順序を構築
        let mut topo_order: Vec<*mut Self> = Vec::new();
        let mut visited = std::collections::HashSet::<usize>::new();
        Self::build_topo(self_ptr, &mut visited, &mut topo_order);

        // 2. 出力テンソルの勾配を ones で初期化（replace, not accumulate）
        let ones = CudaTensor::ones(self.shape(), self.dtype());
        if let Some(ref mut meta) = self.autograd {
            meta.grad = Some(ones);
        }

        // 3. 逆トポロジカル順序で勾配を伝播
        for &ptr in topo_order.iter().rev() {
            let tensor = unsafe { &*ptr };
            let grad = tensor.autograd.as_ref()
                .and_then(|m| m.grad.as_ref())
                .map(|g| g.shallow_clone());

            if let Some(grad_output) = grad {
                let propagation = tensor.autograd.as_ref().and_then(|m| {
                    m.grad_fn.as_ref().map(|gf| {
                        let grads = gf.backward(&grad_output);
                        let inputs = gf.inputs();
                        (grads, inputs)
                    })
                });

                if let Some((grads_result, inputs)) = propagation {
                    let grads = match grads_result {
                        Ok(g) => g,
                        Err(e) => {
                            eprintln!("Autograd backward error: {}", e);
                            continue;
                        }
                    };
                    for (input_ref, grad) in inputs.into_iter().zip(grads.into_iter()) {
                        let input = unsafe { &mut *input_ref.get() };
                        if input.requires_grad() {
                            input.accumulate_grad(grad)?;
                        }
                    }
                }
            }
        }

        // 4. 計算グラフ解放
        // 中間ノード（grad_fn を持つノード）の grad も不要なのでクリアする。
        // リーフテンソル（grad_fn == None）の grad は .grad() で取得するため保持する。
        crate::stream::sync_stream();
        for &ptr in topo_order.iter() {
            let tensor = unsafe { &mut *ptr };
            if let Some(ref mut meta) = tensor.autograd {
                let is_intermediate = meta.grad_fn.is_some();
                meta.grad_fn = None;
                // 中間ノードの勾配はbackward後に不要 → メモリ解放
                if is_intermediate {
                    meta.grad = None;
                }
            }
        }
        Ok(())
    }

    /// DFS でトポロジカル順序を構築（Metal 版と同一ロジック）
    fn build_topo(
        ptr: *mut Self,
        visited: &mut std::collections::HashSet<usize>,
        topo: &mut Vec<*mut Self>,
    ) {
        let key = ptr as usize;
        if visited.contains(&key) {
            return;
        }
        visited.insert(key);
        let tensor = unsafe { &*ptr };
        if let Some(ref meta) = tensor.autograd {
            if let Some(ref gf) = meta.grad_fn {
                for input_ref in gf.inputs() {
                    let input_ptr = input_ref.get() as *mut Self;
                    let input_tensor = unsafe { &*input_ptr };
                    if input_tensor.requires_grad() {
                        Self::build_topo(input_ptr, visited, topo);
                    }
                }
            }
        }
        topo.push(ptr);
    }

    /// 計算グラフから切り離す
    pub fn detach(&self) -> CudaTensor {
        self.shallow_clone()
    }
}

impl Drop for CudaTensor {
    fn drop(&mut self) {
        // バッファの唯一の所有者なら、プールに返却して再利用
        if Arc::strong_count(&self.buffer) == 1 {
            pool_release(self.buffer.clone());
        }
    }
}
