quickraw 0.1.6

use super::byte_stream::ByteStream;
use super::huffman::HuffTable;
use super::*;
use decompressors::*;

mod decompressors;

#[allow(clippy::upper_case_acronyms)]
enum Marker {
    Stuff = 0x00,
    SOF3 = 0xc3, // lossless
    DHT = 0xc4,  // huffman tables
    SOI = 0xd8,  // start of image
    EOI = 0xd9,  // end of image
    SOS = 0xda,  // start of scan
    DQT = 0xdb,  // quantization tables
    Fill = 0xff,
}

fn m(marker: Marker) -> u8 {
    marker as u8
}

#[derive(Debug, Copy, Clone)]
struct JpegComponentInfo {
    // These values are fixed over the whole image, read from the SOF marker.
    id: usize,    // identifier for this component (0..255)
    // index: usize, // its index in SOF or cPtr->compInfo[]

    // Huffman table selector (0..3). The value may vary between scans.
    // It is read from the SOS marker.
    dc_tbl_num: usize,
    super_h: usize, // Horizontal Supersampling
    super_v: usize, // Vertical Supersampling
}

#[derive(Debug, Clone)]
struct SOFInfo {
    width: usize,
    height: usize,
    cps: usize,
    precision: usize,
    components: Vec<JpegComponentInfo>,
    csfix: bool,
}

impl SOFInfo {
    fn empty(csfix: bool) -> SOFInfo {
        SOFInfo {
            width: 0,
            height: 0,
            cps: 0,
            precision: 0,
            components: Vec::new(),
            csfix,
        }
    }

    fn parse_sof(&mut self, input: &mut ByteStream) -> Result<(), DecodingError> {
        let header_length = input.get_u16() as usize;
        self.precision = input.get_u8() as usize;
        self.height = input.get_u16() as usize;
        self.width = input.get_u16() as usize;
        self.cps = input.get_u8() as usize;

        if self.precision > 16 {
            return Err(DecodingError::LJpegError("ljpeg: More than 16 bits per channel is not supported."));
        }
        if self.cps > 4 || self.cps < 1 {
            return Err(DecodingError::LJpegError("ljpeg: Only from 1 to 4 components are supported."));
        }
        if header_length != 8 + self.cps * 3 {
            return Err(DecodingError::LJpegError("ljpeg: Header size mismatch."));
        }

        for _ in 0..self.cps {
            let id = input.get_u8() as usize;
            let subs = input.get_u8() as usize;
            input.get_u8(); // Skip info about quantized

            self.components.push(JpegComponentInfo {
                id,
                // index: i,
                dc_tbl_num: 0,
                super_v: subs & 0xf,
                super_h: subs >> 4,
            });
        }
        Ok(())
    }

    fn parse_sos(&mut self, input: &mut ByteStream) -> Result<(usize, usize), DecodingError> {
        if self.width == 0 {
            return Err(DecodingError::LJpegError("ljpeg: Trying to parse SOS before SOF"));
        }
        input.get_u16(); //skip header length
        let soscps = input.get_u8() as usize;
        if self.cps != soscps {
            return Err(DecodingError::LJpegError("ljpeg: component number mismatch in SOS"));
        }
        for cs in 0..self.cps {
            // At least some MOS cameras have this broken
            let readcs = input.get_u8() as usize;
            let cs = if self.csfix { cs } else { readcs };
            let component = match self.components.iter_mut().find(|&&mut c| c.id == cs) {
                Some(val) => val,
                None => return Err(DecodingError::LJpegError("ljpeg: invalid component selector")),
            };
            let td = (input.get_u8() as usize) >> 4;
            if td > 3 {
                return Err(DecodingError::LJpegError("ljpeg: Invalid Huffman table selection"));
            }
            component.dc_tbl_num = td;
        }
        let pred = input.get_u8() as usize;
        input.get_u8(); // Se + Ah Not used in LJPEG
        let pt = (input.get_u8() as usize) & 0xf; // Point Transform
        Ok((pred, pt))
    }
}

#[derive(Debug)]
pub(in super::super) struct LjpegDecompressor<'a> {
    buffer: &'a [u8],
    sof: SOFInfo,
    predictor: usize,
    point_transform: usize,
    dhts: Vec<HuffTable>,
}

impl<'a> LjpegDecompressor<'a> {
    pub(in super::super) fn new(src: &'a [u8]) -> Result<LjpegDecompressor, DecodingError> {
        LjpegDecompressor::new_full(src, false, false)
    }

    pub(in super::super) fn new_full(src: &'a [u8], dng_bug: bool, csfix: bool) -> Result<LjpegDecompressor, DecodingError> {
        let mut input = ByteStream::new(src, false);
        if LjpegDecompressor::get_next_marker(&mut input, false)? != m(Marker::SOI) {
            return Err(DecodingError::LJpegError("ljpeg: Image did not start with SOI. Probably not LJPEG"));
        }

        let mut sof = SOFInfo::empty(csfix);
        let mut dht_init = [false; 4];
        let mut dht_bits = [[0u32; 17]; 4];
        let mut dht_huffval = [[0u32; 256]; 4];
        let pred;
        let pt;
        loop {
            let marker = LjpegDecompressor::get_next_marker(&mut input, true)?;
            if marker == m(Marker::SOF3) {
                // Start of the frame, giving us the basic info
                sof.parse_sof(&mut input)?;
                if sof.precision > 16 || sof.precision < 12 {
                    return Err(DecodingError::LJpegError("sof.precision"));
                }
            } else if marker == m(Marker::DHT) {
                // Huffman table settings
                LjpegDecompressor::parse_dht(&mut input, &mut dht_init, &mut dht_bits, &mut dht_huffval)?;
            } else if marker == m(Marker::SOS) {
                // Start of the actual stream, we can decode after this
                let (a, b) = sof.parse_sos(&mut input)?;
                pred = a;
                pt = b;
                break;
            } else if marker == m(Marker::EOI) {
                // Should never be reached as we stop at SOS
                return Err(DecodingError::LJpegError("ljpeg: reached EOI before SOS"));
            } else if marker == m(Marker::DQT) {
                return Err(DecodingError::LJpegError("ljpeg: not a valid raw file, found DQT"));
            }
        }

        let mut dhts = Vec::new();
        for i in 0..4 {
            dhts.push(if dht_init[i] {
                HuffTable::new(dht_bits[i], dht_huffval[i], dng_bug)
            } else {
                HuffTable::empty()
            });
        }

        let offset = input.get_pos();
        Ok(LjpegDecompressor {
            buffer: &src[offset..],
            sof,
            predictor: pred,
            point_transform: pt,
            dhts,
        })
    }

    fn get_next_marker(input: &mut ByteStream, allowskip: bool) -> Result<u8, DecodingError> {
        if !allowskip {
            if input.get_u8() != 0xff {
                return Err(DecodingError::LJpegError("ljpeg: (noskip) expected marker not found"));
            }
            let mark = input.get_u8();
            if mark == m(Marker::Stuff) || mark == m(Marker::Fill) {
                return Err(DecodingError::LJpegError("ljpeg: (noskip) expected marker but found stuff or fill"));
            }
            return Ok(mark);
        }
        input.skip_to_marker()?;

        Ok(input.get_u8())
    }

    fn parse_dht(
        input: &mut ByteStream,
        init: &mut [bool; 4],
        bits: &mut [[u32; 17]; 4],
        huffval: &mut [[u32; 256]; 4],
    ) -> Result<(), DecodingError> {
        let mut length = (input.get_u16() as usize) - 2;

        while length > 0 {
            let b = input.get_u8() as usize;
            let tc = b >> 4;
            let th = b & 0xf;

            if tc != 0 {
                return Err(DecodingError::LJpegError("ljpeg: unsupported table class in DHT"));
            }
            if th > 3 {
                return Err(DecodingError::LJpegError("ljpeg: unsupported table id"));
            }

            let mut acc: usize = 0;
            for i in 0..16 {
                bits[th][i + 1] = input.get_u8() as u32;
                acc += bits[th][i + 1] as usize;
            }
            bits[th][0] = 0;

            if acc > 256 {
                return Err(DecodingError::LJpegError("ljpeg: invalid DHT table"));
            }

            if length < 1 + 16 + acc {
                return Err(DecodingError::LJpegError("ljpeg: invalid DHT table length"));
            }

            for i in 0..acc {
                huffval[th][i] = input.get_u8() as u32;
            }

            init[th] = true;
            length -= 1 + 16 + acc;
        }

        Ok(())
    }

    pub(in super::super) fn decode(
        &self,
        out: &mut [u16],
        x: usize,
        stripwidth: usize,
        width: usize,
        height: usize,
        dummy: bool,
    ) -> Result<(), DecodingError> {
        if dummy {
            return Ok(());
        }

        if self.sof.components[0].super_h == 2 && self.sof.components[0].super_v == 2 {
            return decode_ljpeg_420(self, out, width, height);
        } else if self.sof.components[0].super_h == 2 && self.sof.components[0].super_v == 1 {
            return decode_ljpeg_422(self, out, width, height);
        }

        match self.predictor {
            1 => match self.sof.cps {
                2 => decode_ljpeg_2components(self, out, x, stripwidth, width, height),
                3 => decode_ljpeg_3components(self, out, x, stripwidth, width, height),
                4 => decode_ljpeg_4components(self, out, width, height),
                c => Err(DecodingError::LJpegComponentFilesNotSupported(c)),
            },
            8 => decode_hasselblad(self, out, width),
            p => Err(DecodingError::LJpegPredictorNotSupported(p)),
        }
    }

    // pub(in super::super) fn decode_leaf(&self, width: usize, height: usize) -> Result<Vec<u16>> {
    //     let mut offsets = vec![0 as usize; 1];
    //     let mut input = ByteStream::new(self.buffer, false);
    //     loop {
    //         match LjpegDecompressor::get_next_marker(&mut input, true) {
    //             Ok(marker) => {
    //                 if marker == m(Marker::EOI) {
    //                     break;
    //                 }
    //                 offsets.push(input.get_pos());
    //             }
    //             Err(_) => break,
    //         }
    //     }
    //     let nstrips = (height - 1) / 8 + 1;
    //     if offsets.len() != nstrips {
    //         return Err(DecodingError::LJpegError("MOS: expecting strips found"));
    //     }

    //     let ref htable1 = self.dhts[self.sof.components[0].dc_tbl_num];
    //     let ref htable2 = self.dhts[self.sof.components[1].dc_tbl_num];
    //     let bpred = 1 << (self.sof.precision - self.point_transform - 1);

    //     let mut out: Vec<u16> = vec![0u16;width * height];
    //     out.chunks_mut(width * 8).enumerate().for_each(|(row, strip)| {
    //       let block = row * 8;
          
    //       let block = block / 8;
    //       let offset = offsets[block];
    //       let nlines = strip.len() / width;
    //       decode_leaf_strip(&self.buffer[offset..], strip, width, nlines, htable1, htable2, bpred).unwrap();
    //     });

    //     Ok(out)
    // }

    // pub(in super::super) fn width(&self) -> usize {
    //     self.sof.width * self.sof.cps
    // }
    // pub(in super::super) fn height(&self) -> usize {
    //     self.sof.height
    // }
    // pub(in super::super) fn super_v(&self) -> usize {
    //     self.sof.components[0].super_v
    // }
    // pub(in super::super) fn super_h(&self) -> usize {
    //     self.sof.components[0].super_h
    // }
}