//! Checkpoint save/load. Dumps the model's `params()` (in their stable order) to
//! a flat binary file and reloads them into a model with matching architecture.
//!
//! Format (little-endian):
//! ```text
//!   magic   : u32  = 0x58545254  ("XTRT")
//!   version : u32  = 1
//!   n_params: u32
//!   repeat n_params times:
//!     ndim  : u32
//!     dims  : [u32; ndim]
//!     data  : [f32; prod(dims)]
//! ```
//! Architecture/config is NOT stored here — the caller rebuilds the model from
//! the same `Config` and `load_into`s the params (the round-trip and resume both
//! know their config). Gated behind `not(no_cuda)` (it round-trips GPU tensors).

#![cfg(not(no_cuda))]

use std::fs::File;
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::Path;
use xtrain_autodiff::tape::Var;
use xtrain_tensor::{Device, Tensor};

const MAGIC: u32 = 0x5854_5254;
const VERSION: u32 = 1;

/// Write every parameter (value) to `path` in `params()` order.
pub fn save(path: &Path, params: &[Var]) -> std::io::Result<()> {
    let mut w = BufWriter::new(File::create(path)?);
    w.write_all(&MAGIC.to_le_bytes())?;
    w.write_all(&VERSION.to_le_bytes())?;
    w.write_all(&(params.len() as u32).to_le_bytes())?;
    for p in params {
        let v = p.value().to_device(Device::Cpu);
        let shape = v.shape();
        w.write_all(&(shape.len() as u32).to_le_bytes())?;
        for &d in shape {
            w.write_all(&(d as u32).to_le_bytes())?;
        }
        for &x in v.as_slice::<f32>() {
            w.write_all(&x.to_le_bytes())?;
        }
    }
    w.flush()
}

/// Read a checkpoint and overwrite each parameter's value in `params` (in order).
/// Shapes must match the saved ones. Tensors are placed on each param's device.
pub fn load_into(path: &Path, params: &[Var]) -> std::io::Result<()> {
    let mut r = BufReader::new(File::open(path)?);
    assert_eq!(read_u32(&mut r)?, MAGIC, "bad checkpoint magic");
    assert_eq!(read_u32(&mut r)?, VERSION, "unsupported checkpoint version");
    let n = read_u32(&mut r)? as usize;
    assert_eq!(n, params.len(), "checkpoint param count != model");

    for p in params {
        let ndim = read_u32(&mut r)? as usize;
        let mut dims = Vec::with_capacity(ndim);
        for _ in 0..ndim {
            dims.push(read_u32(&mut r)? as usize);
        }
        let numel: usize = dims.iter().product();
        let mut data = vec![0.0f32; numel];
        for slot in data.iter_mut() {
            *slot = read_f32(&mut r)?;
        }
        let device = p.value().device();
        assert_eq!(
            p.value().shape(),
            dims.as_slice(),
            "checkpoint shape mismatch"
        );
        p.set_value(Tensor::from_slice(&data, &dims).to_device(device));
    }
    Ok(())
}

fn read_u32<R: Read>(r: &mut R) -> std::io::Result<u32> {
    let mut b = [0u8; 4];
    r.read_exact(&mut b)?;
    Ok(u32::from_le_bytes(b))
}

fn read_f32<R: Read>(r: &mut R) -> std::io::Result<f32> {
    let mut b = [0u8; 4];
    r.read_exact(&mut b)?;
    Ok(f32::from_le_bytes(b))
}