/// CNEF container format — lossless compressed NEF.
///
/// The file is split into ordered segments at chunk boundaries. Each segment
/// is compressed independently. Decompression concatenates the segments in
/// order, reconstructing the original NEF byte-for-byte.
///
/// ```text
/// Header:
///   magic: [u8; 4] = "CNEF"
///   version: u8 = 1
///   original_file_size: u64 LE
///   segment_count: u32 LE
///
/// Per segment (header then payload, in sequence):
///   type: u8
///   original_offset: u64 LE
///   original_length: u64 LE
///   compressed_length: u64 LE
///   -- if type == RAW_PIXELS (2): raw strip metadata follows --
///     width: u32 LE
///     height: u32 LE
///     bits_per_sample: u8
///     huff_select: u8
///     split_row: u32 LE
///     initial_predictors: [i32; 4] LE  (pUp[0][0], pUp[0][1], pUp[1][0], pUp[1][1])
///   [payload: compressed_length bytes]
/// ```

use std::io::{Read, Seek, SeekFrom, Write};

use crate::jxl;
use crate::nef::{NefChunks, NikonLosslessMeta};
use crate::nikon_lossless;
use crate::tiff::COMPRESSION_NIKON_LOSSLESS;

const MAGIC: &[u8; 4] = b"CNEF";
const VERSION: u8 = 1;

const ZSTD_LEVEL: i32 = 19;

#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SegmentType {
    Zstd = 1,
    RawPixelsJxl = 2,
    JpegJxl = 3,
}

impl SegmentType {
    fn from_u8(v: u8) -> Result<Self, String> {
        match v {
            1 => Ok(Self::Zstd),
            2 => Ok(Self::RawPixelsJxl),
            3 => Ok(Self::JpegJxl),
            _ => Err(format!("unknown segment type {v}")),
        }
    }
}

struct Segment {
    seg_type: SegmentType,
    original_offset: u64,
    original_length: u64,
    payload: Vec<u8>,
    raw_meta: Option<RawPixelsMeta>,
}

#[derive(Clone)]
struct RawPixelsMeta {
    width: u32,
    height: u32,
    bits_per_sample: u8,
    huff_select: u8,
    split_row: u32,
    initial_predictors: [i32; 4],
}

pub fn compress<R: Read + Seek, W: Write>(
    nef: &mut R,
    chunks: &NefChunks,
    lossless_meta: Option<&NikonLosslessMeta>,
    effort: i64,
    out: &mut W,
) -> Result<CompressionStats, String> {
    let mut regions: Vec<(u64, u64, RegionKind)> = Vec::new();

    // Collect all chunk regions
    if chunks.raw_strip.compression == COMPRESSION_NIKON_LOSSLESS && lossless_meta.is_some() {
        regions.push((
            chunks.raw_strip.offset,
            chunks.raw_strip.length,
            RegionKind::RawStrip,
        ));
    } else {
        regions.push((
            chunks.raw_strip.offset,
            chunks.raw_strip.length,
            RegionKind::Blob,
        ));
    }
    for jpeg in &chunks.jpegs {
        regions.push((jpeg.offset, jpeg.length as u64, RegionKind::Jpeg));
    }

    regions.sort_by_key(|r| r.0);

    // Build segments: interleave skeleton gaps with chunk regions
    let mut segments: Vec<Segment> = Vec::new();
    let mut pos: u64 = 0;

    for &(offset, length, kind) in &regions {
        if offset > pos {
            let gap = read_range(nef, pos, (offset - pos) as usize)?;
            segments.push(Segment {
                seg_type: SegmentType::Zstd,
                original_offset: pos,
                original_length: offset - pos,
                payload: zstd_compress(&gap)?,
                raw_meta: None,
            });
        }

        match kind {
            RegionKind::RawStrip => {
                let meta = lossless_meta.unwrap();
                let compressed_raw = read_range(nef, offset, length as usize)?;
                let width = chunks.raw_strip.width as usize;
                let height = chunks.raw_strip.height as usize;
                let bps = chunks.raw_strip.bits_per_sample;

                let pixels = nikon_lossless::decode(
                    &compressed_raw,
                    width,
                    height,
                    bps,
                    meta.huff_select,
                    meta.initial_predictors,
                    meta.split_row,
                )?;

                let compressed_pixels = jxl::encode_pixels(
                    &pixels,
                    chunks.raw_strip.width,
                    chunks.raw_strip.height,
                    bps,
                    effort,
                )?;

                segments.push(Segment {
                    seg_type: SegmentType::RawPixelsJxl,
                    original_offset: offset,
                    original_length: length,
                    payload: compressed_pixels,
                    raw_meta: Some(RawPixelsMeta {
                        width: chunks.raw_strip.width,
                        height: chunks.raw_strip.height,
                        bits_per_sample: bps as u8,
                        huff_select: meta.huff_select as u8,
                        split_row: meta.split_row as u32,
                        initial_predictors: [
                            meta.initial_predictors[0][0],
                            meta.initial_predictors[0][1],
                            meta.initial_predictors[1][0],
                            meta.initial_predictors[1][1],
                        ],
                    }),
                });
            }
            RegionKind::Jpeg => {
                let jpeg_data = read_range(nef, offset, length as usize)?;
                let jxl_data = jxl::encode_jpeg(&jpeg_data)?;
                segments.push(Segment {
                    seg_type: SegmentType::JpegJxl,
                    original_offset: offset,
                    original_length: length,
                    payload: jxl_data,
                    raw_meta: None,
                });
            }
            RegionKind::Blob => {
                let data = read_range(nef, offset, length as usize)?;
                segments.push(Segment {
                    seg_type: SegmentType::Zstd,
                    original_offset: offset,
                    original_length: length,
                    payload: zstd_compress(&data)?,
                    raw_meta: None,
                });
            }
        }

        pos = offset + length;
    }

    // Trailing skeleton after last chunk
    if pos < chunks.file_size {
        let gap = read_range(nef, pos, (chunks.file_size - pos) as usize)?;
        segments.push(Segment {
            seg_type: SegmentType::Zstd,
            original_offset: pos,
            original_length: chunks.file_size - pos,
            payload: zstd_compress(&gap)?,
            raw_meta: None,
        });
    }

    // Write CNEF
    let mut total_compressed: u64 = 0;
    let mut total_original: u64 = 0;
    let mut seg_stats: Vec<SegmentStats> = Vec::new();

    // Header
    out.write_all(MAGIC).w()?;
    out.write_all(&[VERSION]).w()?;
    out.write_all(&chunks.file_size.to_le_bytes()).w()?;
    out.write_all(&(segments.len() as u32).to_le_bytes()).w()?;

    for seg in &segments {
        out.write_all(&[seg.seg_type as u8]).w()?;
        out.write_all(&seg.original_offset.to_le_bytes()).w()?;
        out.write_all(&seg.original_length.to_le_bytes()).w()?;
        out.write_all(&(seg.payload.len() as u64).to_le_bytes()).w()?;

        if let Some(ref m) = seg.raw_meta {
            out.write_all(&m.width.to_le_bytes()).w()?;
            out.write_all(&m.height.to_le_bytes()).w()?;
            out.write_all(&[m.bits_per_sample]).w()?;
            out.write_all(&[m.huff_select]).w()?;
            out.write_all(&m.split_row.to_le_bytes()).w()?;
            for &p in &m.initial_predictors {
                out.write_all(&p.to_le_bytes()).w()?;
            }
        }

        out.write_all(&seg.payload).w()?;

        let payload_len = seg.payload.len() as u64;
        seg_stats.push(SegmentStats {
            seg_type: seg.seg_type,
            original_size: seg.original_length,
            compressed_size: payload_len,
        });
        total_compressed += payload_len;
        total_original += seg.original_length;
    }

    Ok(CompressionStats {
        original_size: total_original,
        compressed_size: total_compressed,
        segment_count: segments.len(),
        segments: seg_stats,
    })
}

pub fn decompress<R: Read, W: Write>(input: &mut R, out: &mut W) -> Result<DecompressionStats, String> {
    // Header
    let mut magic = [0u8; 4];
    input.read_exact(&mut magic).r()?;
    if &magic != MAGIC {
        return Err(format!("not a CNEF file (magic: {:?})", magic));
    }
    let mut ver = [0u8; 1];
    input.read_exact(&mut ver).r()?;
    if ver[0] != VERSION {
        return Err(format!("unsupported CNEF version {}", ver[0]));
    }

    let original_file_size = read_u64_le(input)?;
    let segment_count = read_u32_le(input)? as usize;

    let mut reconstructed_size: u64 = 0;

    for _ in 0..segment_count {
        let mut type_buf = [0u8; 1];
        input.read_exact(&mut type_buf).r()?;
        let seg_type = SegmentType::from_u8(type_buf[0])?;

        let _original_offset = read_u64_le(input)?;
        let original_length = read_u64_le(input)?;
        let compressed_length = read_u64_le(input)? as usize;

        let raw_meta = if seg_type == SegmentType::RawPixelsJxl {
            Some(RawPixelsMeta {
                width: read_u32_le(input)?,
                height: read_u32_le(input)?,
                bits_per_sample: read_u8(input)?,
                huff_select: read_u8(input)?,
                split_row: read_u32_le(input)?,
                initial_predictors: [
                    read_i32_le(input)?,
                    read_i32_le(input)?,
                    read_i32_le(input)?,
                    read_i32_le(input)?,
                ],
            })
        } else {
            None
        };

        let mut payload = vec![0u8; compressed_length];
        input.read_exact(&mut payload).r()?;

        match seg_type {
            SegmentType::Zstd => {
                let decompressed = zstd_decompress(&payload, original_length as usize)?;
                out.write_all(&decompressed).w()?;
                reconstructed_size += decompressed.len() as u64;
            }
            SegmentType::JpegJxl => {
                let jpeg_bytes = jxl::decode_jpeg(&payload)?;
                let needed = original_length as usize;
                if jpeg_bytes.len() != needed {
                    return Err(format!(
                        "JPEG reconstruction size {} != expected {}",
                        jpeg_bytes.len(),
                        needed,
                    ));
                }
                out.write_all(&jpeg_bytes).w()?;
                reconstructed_size += needed as u64;
            }
            SegmentType::RawPixelsJxl => {
                let m = raw_meta.unwrap();
                let pixels = jxl::decode_pixels(&payload, m.width, m.height)?;

                let nikon_compressed = nikon_lossless::encode(
                    &pixels,
                    m.width as usize,
                    m.height as usize,
                    m.bits_per_sample as u32,
                    m.huff_select as usize,
                    [
                        [m.initial_predictors[0], m.initial_predictors[1]],
                        [m.initial_predictors[2], m.initial_predictors[3]],
                    ],
                    m.split_row as usize,
                )?;

                // The Nikon encoder may produce fewer bytes than original due
                // to trailing padding in the original strip. Pad to match.
                let needed = original_length as usize;
                if nikon_compressed.len() > needed {
                    return Err(format!(
                        "re-encoded raw strip {} bytes > original {} bytes",
                        nikon_compressed.len(),
                        needed,
                    ));
                }
                out.write_all(&nikon_compressed).w()?;
                if nikon_compressed.len() < needed {
                    let pad = vec![0u8; needed - nikon_compressed.len()];
                    out.write_all(&pad).w()?;
                }
                reconstructed_size += needed as u64;
            }
        }
    }

    if reconstructed_size != original_file_size {
        return Err(format!(
            "reconstructed size {reconstructed_size} != expected {original_file_size}"
        ));
    }

    Ok(DecompressionStats {
        original_size: original_file_size,
    })
}

#[derive(Debug)]
pub struct CompressionStats {
    pub original_size: u64,
    pub compressed_size: u64,
    pub segment_count: usize,
    pub segments: Vec<SegmentStats>,
}

#[derive(Debug)]
pub struct SegmentStats {
    pub seg_type: SegmentType,
    pub original_size: u64,
    pub compressed_size: u64,
}

#[derive(Debug)]
pub struct DecompressionStats {
    pub original_size: u64,
}

#[derive(Clone, Copy)]
enum RegionKind {
    RawStrip,
    Jpeg,
    Blob,
}

fn read_range<R: Read + Seek>(reader: &mut R, offset: u64, length: usize) -> Result<Vec<u8>, String> {
    reader
        .seek(SeekFrom::Start(offset))
        .map_err(|e| format!("seek to {offset}: {e}"))?;
    let mut buf = vec![0u8; length];
    reader
        .read_exact(&mut buf)
        .map_err(|e| format!("read {length} bytes at {offset}: {e}"))?;
    Ok(buf)
}

fn zstd_compress(data: &[u8]) -> Result<Vec<u8>, String> {
    zstd::encode_all(std::io::Cursor::new(data), ZSTD_LEVEL)
        .map_err(|e| format!("zstd compress: {e}"))
}

fn zstd_decompress(data: &[u8], _expected_size: usize) -> Result<Vec<u8>, String> {
    zstd::decode_all(std::io::Cursor::new(data)).map_err(|e| format!("zstd decompress: {e}"))
}

fn read_u8<R: Read>(r: &mut R) -> Result<u8, String> {
    let mut buf = [0u8; 1];
    r.read_exact(&mut buf).r()?;
    Ok(buf[0])
}

fn read_u32_le<R: Read>(r: &mut R) -> Result<u32, String> {
    let mut buf = [0u8; 4];
    r.read_exact(&mut buf).r()?;
    Ok(u32::from_le_bytes(buf))
}

fn read_i32_le<R: Read>(r: &mut R) -> Result<i32, String> {
    let mut buf = [0u8; 4];
    r.read_exact(&mut buf).r()?;
    Ok(i32::from_le_bytes(buf))
}

fn read_u64_le<R: Read>(r: &mut R) -> Result<u64, String> {
    let mut buf = [0u8; 8];
    r.read_exact(&mut buf).r()?;
    Ok(u64::from_le_bytes(buf))
}

trait IoResultExt<T> {
    fn r(self) -> Result<T, String>;
    fn w(self) -> Result<T, String>;
}

impl<T> IoResultExt<T> for Result<T, std::io::Error> {
    fn r(self) -> Result<T, String> {
        self.map_err(|e| format!("read error: {e}"))
    }
    fn w(self) -> Result<T, String> {
        self.map_err(|e| format!("write error: {e}"))
    }
}
