use std::io::BufReader;
use std::io::Read;
use imgref::ImgVec;
use gifski::Collector;
use y4m::{Colorspace, Decoder, ParseError};
use yuv::color::{MatrixCoefficients, Range};
use yuv::convert::RGBConvert;
use yuv::YUV;
use crate::{SrcPath, BinResult};
use crate::source::{Fps, Source};

pub struct Y4MDecoder {
    fps: Fps,
    in_color_space: Option<MatrixCoefficients>,
    decoder: Decoder<Box<BufReader<dyn Read>>>,
    file_size: Option<u64>,
}

impl Y4MDecoder {
    pub fn new(src: SrcPath, fps: Fps, in_color_space: Option<MatrixCoefficients>) -> BinResult<Self> {
        let mut file_size = None;
        let reader = match src {
            SrcPath::Path(path) => {
                let f = std::fs::File::open(path)?;
                let m = f.metadata()?;
                #[cfg(unix)] {
                    use std::os::unix::fs::MetadataExt;
                    file_size = Some(m.size());
                }
                #[cfg(windows)] {
                    use std::os::windows::fs::MetadataExt;
                    file_size = Some(m.file_size());
                }
                Box::new(BufReader::new(f)) as Box<BufReader<dyn Read>>
            },
            SrcPath::Stdin(buf) => Box::new(buf) as Box<BufReader<dyn Read>>,
        };

        Ok(Self {
            file_size,
            fps,
            in_color_space,
            decoder: Decoder::new(reader).map_err(|e| match e {
                y4m::Error::EOF => "The y4m file is truncated or invalid",
                y4m::Error::BadInput => "The y4m file contains invalid metadata",
                y4m::Error::UnknownColorspace => "y4m uses an unusual color format that is not supported",
                y4m::Error::OutOfMemory => "Out of memory, or the y4m file has bogus dimensions",
                y4m::Error::ParseError(ParseError::InvalidY4M) => "The input is not a y4m file",
                y4m::Error::ParseError(error) => return format!("y4m contains invalid data: {error}"),
                y4m::Error::IoError(error) => return format!("I/O error when reading a y4m file: {error}"),
            }.to_string())?,
        })
    }
}

enum Samp {
    Mono,
    S1x1,
    S2x1,
    S2x2,
}

impl Source for Y4MDecoder {
    fn total_frames(&self) -> Option<u64> {
        self.file_size.map(|file_size| {
            let w = self.decoder.get_width();
            let h = self.decoder.get_height();
            let d = self.decoder.get_bytes_per_sample();
            let s = match self.decoder.get_colorspace() {
                Colorspace::Cmono => 4,
                Colorspace::Cmono12 => 4,
                Colorspace::C420 => 6,
                Colorspace::C420p10 => 6,
                Colorspace::C420p12 => 6,
                Colorspace::C420jpeg => 6,
                Colorspace::C420paldv => 6,
                Colorspace::C420mpeg2 => 6,
                Colorspace::C422 => 8,
                Colorspace::C422p10 => 8,
                Colorspace::C422p12 => 8,
                Colorspace::C444 => 12,
                Colorspace::C444p10 => 12,
                Colorspace::C444p12 => 12,
                _ => 12,
            };
            file_size.saturating_sub(self.decoder.get_raw_params().len() as _) / (w * h * d * s / 4 + 6) as u64
        })
    }

    fn collect(&mut self, c: &mut Collector) -> BinResult<()> {
        let fps = self.decoder.get_framerate();
        let frame_time = 1. / (fps.num as f64 / fps.den as f64);
        let wanted_frame_time = 1. / f64::from(self.fps.fps);
        let width = self.decoder.get_width();
        let height = self.decoder.get_height();
        let raw_params_str = &*String::from_utf8_lossy(self.decoder.get_raw_params()).into_owned();
        let range = raw_params_str.split_once("COLORRANGE=").map(|(_, r)| {
            if r.starts_with("FULL") { Range::Full } else { Range::Limited }
        });

        let matrix = self.in_color_space.unwrap_or({
            if height <= 480 && width <= 720 { MatrixCoefficients::BT601 } else { MatrixCoefficients::BT709 }
        });

        let (samp, conv) = match self.decoder.get_colorspace() {
            Colorspace::Cmono => (Samp::Mono, RGBConvert::<u8>::new(range.unwrap_or(Range::Limited), MatrixCoefficients::Identity)),
            Colorspace::Cmono12 => return Err("Y4M with Cmono12 is not supported yet".into()),
            Colorspace::C420 => (Samp::S2x2, RGBConvert::<u8>::new(range.unwrap_or(Range::Limited), matrix)),
            Colorspace::C420p10 => return Err("Y4M with C420p10 is not supported yet".into()),
            Colorspace::C420p12 => return Err("Y4M with C420p12 is not supported yet".into()),
            Colorspace::C420jpeg => (Samp::S2x2, RGBConvert::<u8>::new(range.unwrap_or(Range::Limited), matrix)),
            Colorspace::C420paldv => (Samp::S2x2, RGBConvert::<u8>::new(range.unwrap_or(Range::Limited), matrix)),
            Colorspace::C420mpeg2 => (Samp::S2x2, RGBConvert::<u8>::new(range.unwrap_or(Range::Limited), matrix)),
            Colorspace::C422 => (Samp::S2x1, RGBConvert::<u8>::new(range.unwrap_or(Range::Limited), matrix)),
            Colorspace::C422p10 => return Err("Y4M with C422p10 is not supported yet".into()),
            Colorspace::C422p12 => return Err("Y4M with C422p12 is not supported yet".into()),
            Colorspace::C444 => (Samp::S1x1, RGBConvert::<u8>::new(range.unwrap_or(Range::Limited), matrix)),
            Colorspace::C444p10 => return Err("Y4M with C444p10 is not supported yet".into()),
            Colorspace::C444p12 => return Err("Y4M with C444p12 is not supported yet".into()),
            _ => return Err(format!("Y4M uses unsupported color mode {raw_params_str}").into()),
        };
        let conv = conv?;
        if width == 0 || width > u16::MAX as _ || height == 0 || height > u16::MAX as _ {
            return Err("Video too large".into());
        }

        #[cold]
        fn bad_frame(mode: &str) -> BinResult<()> {
            Err(format!("Bad Y4M frame (using {mode})").into())
        }

        let mut idx = 0;
        let mut presentation_timestamp = 0.0;
        let mut wanted_pts = 0.0;
        loop {
            match self.decoder.read_frame() {
                Ok(frame) => {
                    let this_frame_pts = presentation_timestamp / f64::from(self.fps.speed);
                    presentation_timestamp += frame_time;
                    if presentation_timestamp < wanted_pts {
                        continue; // skip a frame
                    }
                    wanted_pts += wanted_frame_time;

                    let y = frame.get_y_plane();
                    if y.is_empty() {
                        return bad_frame(raw_params_str);
                    }
                    let u = frame.get_u_plane();
                    let v = frame.get_v_plane();
                    if v.len() != u.len() {
                        return bad_frame(raw_params_str);
                    }

                    let mut out = Vec::new();
                    out.try_reserve(width * height)?;
                    match samp {
                        Samp::Mono => todo!(),
                        Samp::S1x1 => {
                            if v.len() != y.len() {
                                return bad_frame(raw_params_str);
                            }

                            let y = y.chunks_exact(width);
                            let u = u.chunks_exact(width);
                            let v = v.chunks_exact(width);
                            if y.len() != v.len() {
                                return bad_frame(raw_params_str);
                            }
                            for (y, (u, v)) in y.zip(u.zip(v)) {
                                out.extend(
                                    y.iter().copied().zip(u.iter().copied().zip(v.iter().copied()))
                                    .map(|(y, (u, v))| {
                                        conv.to_rgb(YUV {y, u, v}).with_alpha(255)
                                    }));
                            }
                        },
                        Samp::S2x1 => {
                            let y = y.chunks_exact(width);
                            let u = u.chunks_exact(width.div_ceil(2));
                            let v = v.chunks_exact(width.div_ceil(2));
                            if y.len() != v.len() {
                                return bad_frame(raw_params_str);
                            }
                            for (y, (u, v)) in y.zip(u.zip(v)) {
                                let u = u.iter().copied().flat_map(|x| [x, x]);
                                let v = v.iter().copied().flat_map(|x| [x, x]);
                                out.extend(
                                    y.iter().copied().zip(u.zip(v))
                                    .map(|(y, (u, v))| {
                                        conv.to_rgb(YUV {y, u, v}).with_alpha(255)
                                    }));
                            }
                        },
                        Samp::S2x2 => {
                            let y = y.chunks_exact(width);
                            let u = u.chunks_exact(width.div_ceil(2)).flat_map(|r| [r, r]);
                            let v = v.chunks_exact(width.div_ceil(2)).flat_map(|r| [r, r]);
                            for (y, (u, v)) in y.zip(u.zip(v)) {
                                let u = u.iter().copied().flat_map(|x| [x, x]);
                                let v = v.iter().copied().flat_map(|x| [x, x]);
                                out.extend(
                                    y.iter().copied().zip(u.zip(v))
                                    .map(|(y, (u, v))| {
                                        conv.to_rgb(YUV {y, u, v}).with_alpha(255)
                                    }));
                            }
                        },
                    }
                    if out.len() != width * height {
                        return bad_frame(raw_params_str);
                    }
                    let pixels = ImgVec::new(out, width, height);

                    c.add_frame_rgba(idx, pixels, this_frame_pts)?;
                    idx += 1;
                },
                Err(y4m::Error::EOF) => break,
                Err(e) => return Err(e.into()),
            }
        }
        Ok(())
    }
}
