use crate::error::{Error, Result};
use crate::jpeg::byteorder::ReadBytesExt;
use crate::jpeg::marker::Marker;
use crate::jpeg::parser::{
    parse_sof, parse_sos, skip_marker, CodingProcess, FrameInfo, ScanInfo,
};
use std::io;
use std::io::Read;

pub struct PositionTrackingReader<R> {
    reader: R,
    position: usize,
}

impl<R> PositionTrackingReader<R> {
    pub fn position(&self) -> usize {
        self.position
    }
}

impl<R: Read> Read for PositionTrackingReader<R> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let n = self.reader.read(buf)?;
        self.position += n;
        Ok(n)
    }

    #[inline]
    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        self.reader.read_exact(buf)?;
        self.position += buf.len();
        Ok(())
    }
}

/// JPEG decoder
pub struct Decoder<R> {
    reader: PositionTrackingReader<R>,
    frame: Option<FrameInfo>,
}

#[derive(Debug)]
pub enum MarkerData {
    Frame(FrameInfo),
    Scan(ScanInfo),
    Other(Marker),
}

#[derive(Debug)]
pub struct MarkerPosition {
    pub position: usize,
    pub marker: MarkerData,
}

impl<R: Read> Decoder<R> {
    /// Creates a new `Decoder` using the reader `reader`.
    pub fn new(reader: R) -> Self {
        Self {
            reader: PositionTrackingReader {
                reader,
                position: 0,
            },
            frame: None,
        }
    }

    pub fn decode(&mut self) -> Result<Vec<MarkerPosition>> {
        let mut positions = Vec::new();
        if self.frame.is_none() && (self.reader.read_u8()? != 0xFF || Marker::from_u8(self.reader.read_u8()?) != Some(Marker::SOI)) {
            return Err(Error::Format("first two bytes is not a SOI marker"));
        }

        let mut pending_marker = None;

        loop {
            let (position, marker) = match pending_marker.take() {
                Some(m) => m,
                None => match self.read_marker() {
                    Ok(m) => m,
                    Err(_) => break,
                },
            };

            match marker {
                // Frame header
                Marker::SOF(..) => {
                    // Section 4.10
                    // "An image contains only one frame in the cases of sequential and
                    //  progressive coding processes; an image contains multiple frames for the
                    //  hierarchical mode."
                    if self.frame.is_some() {
                        return Err(Error::Unsupported);
                    }

                    let frame = parse_sof(&mut self.reader, marker)?;
                    let component_count = frame.components.len();

                    if frame.is_differential {
                        return Err(Error::Unsupported);
                    }
                    if frame.coding_process == CodingProcess::Lossless {
                        return Err(Error::Unsupported);
                    }
                    if frame.precision != 8 {
                        return Err(Error::Unsupported);
                    }
                    if frame.image_size.height == 0 {
                        return Err(Error::Unsupported);
                    }
                    if component_count != 1 && component_count != 3 && component_count != 4 {
                        return Err(Error::Unsupported);
                    }
                    self.frame = Some(frame.clone());

                    positions.push(MarkerPosition {
                        position,
                        marker: MarkerData::Frame(frame),
                    });
                },

                // Scan header
                Marker::SOS => {
                    if self.frame.is_none() {
                        return Err(Error::Format("scan encountered before frame"));
                    }

                    let frame = self.frame.clone().ok_or(Error::Unsupported)?;
                    let scan = parse_sos(&mut self.reader, &frame)?;

                    pending_marker = Some(self.skip_scan()?);

                    positions.push(MarkerPosition {
                        position,
                        marker: MarkerData::Scan(scan),
                    });
                },

                // Restart
                Marker::RST(..) => {
                },

                // End of image
                Marker::EOI => {
                    positions.push(MarkerPosition {
                        position,
                        marker: MarkerData::Other(marker),
                    });
                    break;
                },

                marker => {
                    skip_marker(&mut self.reader, marker)?;

                    positions.push(MarkerPosition {
                        position,
                        marker: MarkerData::Other(marker),
                    });
                },
            }
        }

        Ok(positions)
    }

    fn read_marker(&mut self) -> Result<(usize, Marker)> {
        // This should be an error as the JPEG spec doesn't allow extraneous data between marker segments.
        // libjpeg allows this though and there are images in the wild utilising it, so we are
        // forced to support this behavior.
        // Sony Ericsson P990i is an example of a device which produce this sort of JPEGs.
        while self.reader.read_u8()? != 0xFF {}

        let mut byte = self.reader.read_u8()?;

        // Section B.1.1.2
        // "Any marker may optionally be preceded by any number of fill bytes, which are bytes assigned code X’FF’."
        while byte == 0xFF {
            byte = self.reader.read_u8()?;
        }

        let position = self.reader.position() - 2;

        match byte {
            0x00 => Err(Error::Format("FF 00 found where marker was expected")),
            _    => Ok((position, Marker::from_u8(byte).ok_or(Error::Unsupported)?)),
        }
    }

    fn skip_scan(&mut self) -> Result<(usize, Marker)> {
        loop {
            let byte = self.reader.read_u8()?;

            if byte == 0xFF {
                let mut next_byte = self.reader.read_u8()?;

                // Check for byte stuffing.
                if next_byte != 0x00 {
                    // We seem to have reached the end of entropy-coded data and encountered a
                    // marker. Since we can't put data back into the reader, we have to continue
                    // reading to identify the marker so we can pass it on.

                    // Section B.1.1.2
                    // "Any marker may optionally be preceded by any number of fill bytes, which are bytes assigned code X’FF’."
                    while next_byte == 0xFF {
                        next_byte = self.reader.read_u8()?;
                    }

                    let position = self.reader.position() - 2;
                    return match next_byte {
                        0x00 => Err(Error::Format("FF 00 found where marker was expected")),
                        _    => Ok((position, Marker::from_u8(next_byte).ok_or(Error::Unsupported)?)),
                    };
                }
            }
        }
    }
}