//! NIfTI file writer supporting .nii and .nii.gz formats.

use crate::error::{MedicalIoError, Result};
use crate::nifti::header::{NiftiDataType, NiftiHeader, SpatialUnits, TransformCode};
use crate::volume::Volume;
use byteorder::{LittleEndian, WriteBytesExt};
use flate2::Compression;
use flate2::write::GzEncoder;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;

/// Write a Volume to a NIfTI file.
///
/// The output format is determined by the file extension:
/// - `.nii` - uncompressed NIfTI-1
/// - `.nii.gz` - gzip-compressed NIfTI-1
///
/// # Arguments
/// * `volume` - The volume to write
/// * `path` - Output file path
/// * `datatype` - Output data type (default: Float64)
///
/// # Example
/// ```ignore
/// use rtx_medical_io::nifti::{write_nifti, NiftiDataType};
/// use rtx_medical_io::volume::Volume;
///
/// let volume = Volume::zeros([64, 64, 64]);
/// write_nifti(&volume, "output.nii.gz", NiftiDataType::Float32)?;
/// ```
pub fn write_nifti<P: AsRef<Path>>(
    volume: &Volume,
    path: P,
    datatype: NiftiDataType,
) -> Result<()> {
    let path = path.as_ref();
    let path_str = path.to_string_lossy().to_lowercase();

    let is_gzipped = path_str.ends_with(".nii.gz") || path_str.ends_with(".gz");

    // Build header
    let header = build_header(volume, datatype);

    // Convert data to target type
    let raw_data = convert_from_f64(volume.data(), datatype)?;

    if is_gzipped {
        write_nifti_gz(path, &header, &raw_data)
    } else {
        write_nifti_uncompressed(path, &header, &raw_data)
    }
}

/// Write a gzip-compressed NIfTI file
fn write_nifti_gz<P: AsRef<Path>>(path: P, header: &NiftiHeader, data: &[u8]) -> Result<()> {
    let file = File::create(path)?;
    let writer = BufWriter::new(file);
    let mut encoder = GzEncoder::new(writer, Compression::default());

    write_header(&mut encoder, header)?;
    encoder.write_all(data)?;
    encoder.finish().map_err(|e| {
        MedicalIoError::Compression(format!("Failed to finish gzip compression: {}", e))
    })?;

    Ok(())
}

/// Write an uncompressed NIfTI file
fn write_nifti_uncompressed<P: AsRef<Path>>(
    path: P,
    header: &NiftiHeader,
    data: &[u8],
) -> Result<()> {
    let file = File::create(path)?;
    let mut writer = BufWriter::new(file);

    write_header(&mut writer, header)?;
    writer.write_all(data)?;
    writer.flush()?;

    Ok(())
}

/// Build a NIfTI-1 header from a Volume
fn build_header(volume: &Volume, datatype: NiftiDataType) -> NiftiHeader {
    let [x, y, z] = volume.shape();
    let [dx, dy, dz] = volume.spacing();
    let affine = volume.affine();

    NiftiHeader {
        sizeof_hdr: 348,
        datatype,
        bitpix: (datatype.bytes_per_voxel() * 8) as i16,
        dim: [3, x as i64, y as i64, z as i64, 1, 1, 1, 1],
        intent_p1: 0.0,
        intent_p2: 0.0,
        intent_p3: 0.0,
        intent_code: 0,
        pixdim: [1.0, dx, dy, dz, 1.0, 1.0, 1.0, 1.0],
        vox_offset: 352,
        scl_slope: 1.0,
        scl_inter: 0.0,
        slice_code: 0,
        xyzt_units: SpatialUnits::Millimeters as u8,
        cal_max: 0.0,
        cal_min: 0.0,
        slice_duration: 0.0,
        toffset: 0.0,
        slice_start: 0,
        slice_end: 0,
        descrip: "Created by rtx-medical-io".to_string(),
        aux_file: String::new(),
        qform_code: TransformCode::Unknown,
        sform_code: TransformCode::ScannerAnat,
        quatern_b: 0.0,
        quatern_c: 0.0,
        quatern_d: 0.0,
        qoffset_x: 0.0,
        qoffset_y: 0.0,
        qoffset_z: 0.0,
        srow_x: affine[0],
        srow_y: affine[1],
        srow_z: affine[2],
        intent_name: String::new(),
        magic: *b"n+1\0\0\0\0\0",
        is_nifti2: false,
        little_endian: true,
    }
}

/// Write NIfTI-1 header (348 bytes + 4 bytes padding = 352 bytes)
fn write_header<W: Write>(writer: &mut W, header: &NiftiHeader) -> Result<()> {
    // sizeof_hdr (0-3)
    writer.write_i32::<LittleEndian>(header.sizeof_hdr)?;

    // data_type (unused, 4-13) - 10 bytes
    writer.write_all(&[0u8; 10])?;

    // db_name (unused, 14-31) - 18 bytes
    writer.write_all(&[0u8; 18])?;

    // extents (unused, 32-35)
    writer.write_i32::<LittleEndian>(0)?;

    // session_error (unused, 36-37)
    writer.write_i16::<LittleEndian>(0)?;

    // regular (unused, 38)
    writer.write_all(&[0u8; 1])?;

    // dim_info (unused, 39)
    writer.write_all(&[0u8; 1])?;

    // dim (40-55) - 8 x i16
    for i in 0..8 {
        writer.write_i16::<LittleEndian>(header.dim[i] as i16)?;
    }

    // intent_p1 (56-59)
    writer.write_f32::<LittleEndian>(header.intent_p1 as f32)?;

    // intent_p2 (60-63)
    writer.write_f32::<LittleEndian>(header.intent_p2 as f32)?;

    // intent_p3 (64-67)
    writer.write_f32::<LittleEndian>(header.intent_p3 as f32)?;

    // intent_code (68-69)
    writer.write_i16::<LittleEndian>(header.intent_code)?;

    // datatype (70-71)
    writer.write_i16::<LittleEndian>(header.datatype as i16)?;

    // bitpix (72-73)
    writer.write_i16::<LittleEndian>(header.bitpix)?;

    // slice_start (74-75)
    writer.write_i16::<LittleEndian>(header.slice_start as i16)?;

    // pixdim (76-107) - 8 x f32
    for i in 0..8 {
        writer.write_f32::<LittleEndian>(header.pixdim[i] as f32)?;
    }

    // vox_offset (108-111)
    writer.write_f32::<LittleEndian>(header.vox_offset as f32)?;

    // scl_slope (112-115)
    writer.write_f32::<LittleEndian>(header.scl_slope as f32)?;

    // scl_inter (116-119)
    writer.write_f32::<LittleEndian>(header.scl_inter as f32)?;

    // slice_end (120-121)
    writer.write_i16::<LittleEndian>(header.slice_end as i16)?;

    // slice_code (122)
    writer.write_all(&[header.slice_code])?;

    // xyzt_units (123)
    writer.write_all(&[header.xyzt_units])?;

    // cal_max (124-127)
    writer.write_f32::<LittleEndian>(header.cal_max as f32)?;

    // cal_min (128-131)
    writer.write_f32::<LittleEndian>(header.cal_min as f32)?;

    // slice_duration (132-135)
    writer.write_f32::<LittleEndian>(header.slice_duration as f32)?;

    // toffset (136-139)
    writer.write_f32::<LittleEndian>(header.toffset as f32)?;

    // glmax (unused, 140-143)
    writer.write_i32::<LittleEndian>(0)?;

    // glmin (unused, 144-147)
    writer.write_i32::<LittleEndian>(0)?;

    // descrip (148-227) - 80 bytes
    let descrip_bytes = header.descrip.as_bytes();
    let mut descrip_buf = [0u8; 80];
    let len = descrip_bytes.len().min(80);
    descrip_buf[..len].copy_from_slice(&descrip_bytes[..len]);
    writer.write_all(&descrip_buf)?;

    // aux_file (228-251) - 24 bytes
    let aux_bytes = header.aux_file.as_bytes();
    let mut aux_buf = [0u8; 24];
    let len = aux_bytes.len().min(24);
    aux_buf[..len].copy_from_slice(&aux_bytes[..len]);
    writer.write_all(&aux_buf)?;

    // qform_code (252-253)
    writer.write_i16::<LittleEndian>(header.qform_code as i16)?;

    // sform_code (254-255)
    writer.write_i16::<LittleEndian>(header.sform_code as i16)?;

    // quatern_b (256-259)
    writer.write_f32::<LittleEndian>(header.quatern_b as f32)?;

    // quatern_c (260-263)
    writer.write_f32::<LittleEndian>(header.quatern_c as f32)?;

    // quatern_d (264-267)
    writer.write_f32::<LittleEndian>(header.quatern_d as f32)?;

    // qoffset_x (268-271)
    writer.write_f32::<LittleEndian>(header.qoffset_x as f32)?;

    // qoffset_y (272-275)
    writer.write_f32::<LittleEndian>(header.qoffset_y as f32)?;

    // qoffset_z (276-279)
    writer.write_f32::<LittleEndian>(header.qoffset_z as f32)?;

    // srow_x (280-295) - 4 x f32
    for val in &header.srow_x {
        writer.write_f32::<LittleEndian>(*val as f32)?;
    }

    // srow_y (296-311) - 4 x f32
    for val in &header.srow_y {
        writer.write_f32::<LittleEndian>(*val as f32)?;
    }

    // srow_z (312-327) - 4 x f32
    for val in &header.srow_z {
        writer.write_f32::<LittleEndian>(*val as f32)?;
    }

    // intent_name (328-343) - 16 bytes
    let intent_bytes = header.intent_name.as_bytes();
    let mut intent_buf = [0u8; 16];
    let len = intent_bytes.len().min(16);
    intent_buf[..len].copy_from_slice(&intent_bytes[..len]);
    writer.write_all(&intent_buf)?;

    // magic (344-347) - 4 bytes for NIfTI-1
    writer.write_all(&header.magic[..4])?;

    // Padding to vox_offset (348-351) - 4 bytes
    writer.write_all(&[0u8; 4])?;

    Ok(())
}

/// Convert f64 data to raw bytes for the target data type
fn convert_from_f64(data: &[f64], datatype: NiftiDataType) -> Result<Vec<u8>> {
    let mut result = Vec::with_capacity(data.len() * datatype.bytes_per_voxel());

    match datatype {
        NiftiDataType::UInt8 => {
            for &val in data {
                result.push(val.clamp(0.0, 255.0) as u8);
            }
        }

        NiftiDataType::Int8 => {
            for &val in data {
                result.push(val.clamp(-128.0, 127.0) as i8 as u8);
            }
        }

        NiftiDataType::Int16 => {
            for &val in data {
                let v = val.clamp(i16::MIN as f64, i16::MAX as f64) as i16;
                result.write_i16::<LittleEndian>(v)?;
            }
        }

        NiftiDataType::UInt16 => {
            for &val in data {
                let v = val.clamp(0.0, u16::MAX as f64) as u16;
                result.write_u16::<LittleEndian>(v)?;
            }
        }

        NiftiDataType::Int32 => {
            for &val in data {
                let v = val.clamp(i32::MIN as f64, i32::MAX as f64) as i32;
                result.write_i32::<LittleEndian>(v)?;
            }
        }

        NiftiDataType::UInt32 => {
            for &val in data {
                let v = val.clamp(0.0, u32::MAX as f64) as u32;
                result.write_u32::<LittleEndian>(v)?;
            }
        }

        NiftiDataType::Float32 => {
            for &val in data {
                result.write_f32::<LittleEndian>(val as f32)?;
            }
        }

        NiftiDataType::Float64 => {
            for &val in data {
                result.write_f64::<LittleEndian>(val)?;
            }
        }

        NiftiDataType::Int64 => {
            for &val in data {
                let v = val.clamp(i64::MIN as f64, i64::MAX as f64) as i64;
                result.write_i64::<LittleEndian>(v)?;
            }
        }

        NiftiDataType::UInt64 => {
            for &val in data {
                let v = val.clamp(0.0, u64::MAX as f64) as u64;
                result.write_u64::<LittleEndian>(v)?;
            }
        }

        _ => {
            return Err(MedicalIoError::UnsupportedDataType(datatype as i16));
        }
    }

    Ok(result)
}

/// Write a Volume to an INR file (INRIMAGE-4 format).
///
/// This format is used internally by CGAL for mesh generation.
/// It's a simple headerless format with a text header followed by raw binary data.
pub fn write_inr<P: AsRef<Path>>(volume: &Volume, path: P) -> Result<()> {
    let path = path.as_ref();
    let file = File::create(path)?;
    let mut writer = BufWriter::new(file);

    let [xdim, ydim, zdim] = volume.shape();
    let [vx, vy, vz] = volume.spacing();

    // Determine data type and bit length
    let (btype, bitlen) = ("float", 64);

    // Build header
    let mut header = format!(
        "#INRIMAGE-4#{{\n\
         XDIM={}\n\
         YDIM={}\n\
         ZDIM={}\n\
         VDIM=1\n\
         TYPE={}\n\
         PIXSIZE={} bits\n\
         CPU=decm\n\
         VX={:.6}\n\
         VY={:.6}\n\
         VZ={:.6}\n",
        xdim, ydim, zdim, btype, bitlen, vx, vy, vz
    );

    // Pad header to 256 - 4 bytes, then add closing tag
    let target_len = 256 - 4;
    while header.len() < target_len {
        header.push('\n');
    }
    header.push_str("##}\n");

    // Write header
    writer.write_all(header.as_bytes())?;

    // Write data in Fortran order (column-major)
    let data = volume.data();
    for &val in data {
        writer.write_f64::<LittleEndian>(val)?;
    }

    writer.flush()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_build_header() {
        let volume = Volume::zeros([64, 64, 32]);
        let header = build_header(&volume, NiftiDataType::Float32);

        assert_eq!(header.sizeof_hdr, 348);
        assert_eq!(header.dim[1], 64);
        assert_eq!(header.dim[2], 64);
        assert_eq!(header.dim[3], 32);
        assert_eq!(header.datatype, NiftiDataType::Float32);
        assert_eq!(header.bitpix, 32);
    }

    #[test]
    fn test_convert_uint8() {
        let data = vec![0.0, 127.5, 255.0, 300.0, -10.0];
        let result = convert_from_f64(&data, NiftiDataType::UInt8).unwrap();
        // 127.5 truncates to 127 (not rounded)
        assert_eq!(result, vec![0, 127, 255, 255, 0]);
    }
}