[nihav-player.git] / hwdec-vaapi / src / lib.rs

use std::collections::VecDeque;
use std::convert::TryInto;
use std::rc::Rc;

use nihav_core::codecs::*;
use nihav_core::io::byteio::*;
use nihav_core::io::bitreader::*;
use nihav_core::io::intcode::*;

use libva::*;

#[cfg(debug_assertions)]
macro_rules! validate {
    ($a:expr) => { if !$a { println!("check failed at {}:{}", file!(), line!()); return Err(DecoderError::InvalidData); } };
}
#[cfg(not(debug_assertions))]
macro_rules! validate {
    ($a:expr) => { if !$a { return Err(DecoderError::InvalidData); } };
}

mod pic_ref;
pub use pic_ref::*;
#[allow(clippy::manual_range_contains)]
#[allow(clippy::needless_range_loop)]
mod sets;
use sets::*;
#[allow(clippy::manual_range_contains)]
mod slice;
use slice::*;

trait ReadUE {
    fn read_ue(&mut self) -> DecoderResult<u32>;
    fn read_te(&mut self, range: u32) -> DecoderResult<u32>;
    fn read_ue_lim(&mut self, max_val: u32) -> DecoderResult<u32> {
        let val = self.read_ue()?;
        validate!(val <= max_val);
        Ok(val)
    }
    fn read_se(&mut self) -> DecoderResult<i32> {
        let val = self.read_ue()?;
        if (val & 1) != 0 {
            Ok (((val >> 1) as i32) + 1)
        } else {
            Ok (-((val >> 1) as i32))
        }
    }
}

impl<'a> ReadUE for BitReader<'a> {
    fn read_ue(&mut self) -> DecoderResult<u32> {
        Ok(self.read_code(UintCodeType::GammaP)? - 1)
    }
    fn read_te(&mut self, range: u32) -> DecoderResult<u32> {
        if range == 1 {
            if self.read_bool()? {
                Ok(0)
            } else {
                Ok(1)
            }
        } else {
            let val = self.read_ue()?;
            validate!(val <= range);
            Ok(val)
        }
    }
}

fn get_long_term_id(is_idr: bool, slice_hdr: &SliceHeader) -> Option<usize> {
    if is_idr && !slice_hdr.long_term_reference {
        None
    } else {
        let marking = &slice_hdr.adaptive_ref_pic_marking;
        for (&op, &arg) in marking.memory_management_control_op.iter().zip(marking.operation_arg.iter()).take(marking.num_ops) {
            if op == 6 {
                return Some(arg as usize);
            }
        }
        None
    }
}

fn unescape_nal(src: &[u8], dst: &mut Vec<u8>) -> usize {
    let mut off = 0;
    let mut zrun = 0;
    dst.clear();
    dst.reserve(src.len());
    while off < src.len() {
        dst.push(src[off]);
        if src[off] != 0 {
            zrun = 0;
        } else {
            zrun += 1;
            if zrun == 2 && off + 1 < src.len() && src[off + 1] == 0x03 {
                zrun = 0;
                off += 1;
            }
            if zrun >= 3 && off + 1 < src.len() && src[off + 1] == 0x01 {
                off -= 3;
                dst.truncate(off);
                break;
            }
        }
        off += 1;
    }
    off
}

fn make_dummy_h264_pic() -> PictureH264 {
    PictureH264::new(VA_INVALID_ID, 0, H264PictureFlag::Invalid.into(), 0, 0)
}

trait MakePicH264 {
    fn make_pic(&self) -> PictureH264;
}

impl MakePicH264 for PictureInfo {
    fn make_pic(&self) -> PictureH264 {
        let mut flags = H264PictureFlags::default();
        let frame_idx = if let Some(id) = self.long_term {
                flags |= H264PictureFlag::LongTermReference;
                id as u32
            } else {
                if self.is_ref {
                    flags |= H264PictureFlag::ShortTermReference;
                }
                u32::from(self.id)
            };
        PictureH264::new(self.surface_id, frame_idx, flags, self.top_id as i32, self.bot_id as i32)
    }
}

fn map_ref_list(refs: &[Option<PictureInfo>]) -> [PictureH264; 32] {
    let mut ref_list = Vec::with_capacity(32);

    for rpic in refs.iter() {
        ref_list.push(rpic.as_ref().map_or_else(make_dummy_h264_pic, |pic| pic.make_pic()));
    }

    while ref_list.len() < 32 {
        ref_list.push(make_dummy_h264_pic());
    }
    if let Ok(ret) = ref_list.try_into() {
        ret
    } else {
        panic!("can't convert");
    }
}

fn profile_name(profile: VAProfile::Type) -> &'static str {
    match profile {
        VAProfile::VAProfileMPEG2Simple => "MPEG2 Simple",
        VAProfile::VAProfileMPEG2Main => "MPEG2 Main",
        VAProfile::VAProfileMPEG4Simple => "MPEG4 Simple",
        VAProfile::VAProfileMPEG4AdvancedSimple => "MPEG4 Advanced Simple",
        VAProfile::VAProfileMPEG4Main => "MPEG4 Main",
        VAProfile::VAProfileH264Baseline => "H264 Baseline",
        VAProfile::VAProfileH264Main => "H264 Main",
        VAProfile::VAProfileH264High => "H264 High",
        VAProfile::VAProfileVC1Simple => "VC1 Simple",
        VAProfile::VAProfileVC1Main => "VC1 Main",
        VAProfile::VAProfileVC1Advanced => "VC1 Advanced",
        VAProfile::VAProfileH263Baseline => "H263 Baseline",
        VAProfile::VAProfileJPEGBaseline => "JPEG Baseline",
        VAProfile::VAProfileH264ConstrainedBaseline => "H264 Constrained Baseline",
        VAProfile::VAProfileVP8Version0_3 => "VP8",
        VAProfile::VAProfileH264MultiviewHigh => "H.264 Multiview High",
        VAProfile::VAProfileH264StereoHigh => "H264 Stereo High",
        VAProfile::VAProfileHEVCMain => "H.EVC Main",
        VAProfile::VAProfileHEVCMain10 => "H.EVC Main10",
        VAProfile::VAProfileVP9Profile0 => "VP9 Profile 0",
        VAProfile::VAProfileVP9Profile1 => "VP9 Profile 1",
        VAProfile::VAProfileVP9Profile2 => "VP9 Profile 2",
        VAProfile::VAProfileVP9Profile3 => "VP9 Profile 3",
        VAProfile::VAProfileHEVCMain12 => "HEVC Main12",
        VAProfile::VAProfileHEVCMain422_10 => "HEVC Main10 4:2:2",
        VAProfile::VAProfileHEVCMain422_12 => "HEVC Main12 4:2:2",
        VAProfile::VAProfileHEVCMain444 => "HEVC Main 4:4:4",
        VAProfile::VAProfileHEVCMain444_10 => "HEVC Main10 4:4:4",
        VAProfile::VAProfileHEVCMain444_12 => "HEVC Main12 4:4:4",
        VAProfile::VAProfileHEVCSccMain => "HEVC SCC Main",
        VAProfile::VAProfileHEVCSccMain10 => "HEVC SCC Main10",
        VAProfile::VAProfileHEVCSccMain444 => "HEVC SCC Main 4:4:4",
        VAProfile::VAProfileAV1Profile0 => "AV1 Profile 0",
        VAProfile::VAProfileAV1Profile1 => "AV1 Profile 1",
        VAProfile::VAProfileHEVCSccMain444_10 => "HEVC SCC Main10 4:4:4",
        _ => "unknown",
    }
}

const NUM_REF_PICS: usize = 16;

struct WaitingFrame {
    ts:     u64,
    pic:    Picture<PictureEnd>,
    is_idr: bool,
    is_ref: bool,
    ftype:  FrameType,
}

struct Reorderer {
    last_ref_dts:   Option<u64>,
    ready_idx:      usize,
    frames:         VecDeque<WaitingFrame>,
}

impl Default for Reorderer {
    fn default() -> Self {
        Self {
            last_ref_dts:   None,
            ready_idx:      0,
            frames:         VecDeque::with_capacity(16),
        }
    }
}

impl Reorderer {
    fn add_frame(&mut self, new_frame: WaitingFrame) {
        if !new_frame.is_ref {
            if self.frames.is_empty() {
                self.frames.push_back(new_frame);
            } else {
                let new_dts = new_frame.ts;
                let mut idx = 0;
                for (i, frm) in self.frames.iter().enumerate() {
                    idx = i;
                    if frm.ts > new_dts {
                        break;
                    }
                }
                self.frames.insert(idx, new_frame);
            }
        } else {
            for (i, frm) in self.frames.iter().enumerate() {
                if Some(frm.ts) == self.last_ref_dts {
                    self.ready_idx = i + 1;
                }
            }
            self.last_ref_dts = Some(new_frame.ts);
            self.frames.push_back(new_frame);
        }
    }
    fn get_frame(&mut self) -> Option<WaitingFrame> {
        if self.ready_idx > 0 {
            match self.frames[0].pic.query_status() {
                _ if self.ready_idx > 16 => {},
                Ok(VASurfaceStatus::Ready) => {},
                Ok(VASurfaceStatus::Rendering) => return None,
                _ => {
                    unimplemented!();
                },
            };
            self.ready_idx -= 1;
            self.frames.pop_front()
        } else {
            None
        }
    }
    fn flush(&mut self) {
        self.last_ref_dts = None;
        self.ready_idx = 0;
    }
}

#[allow(dead_code)]
struct VaapiInternals {
    display:        Rc<Display>,
    context:        Rc<Context>,
    ref_pics:       Vec<(Picture<PictureSync>, VASurfaceID)>,
    surfaces:       Vec<Surface>,
    ifmt:           VAImageFormat,
}

pub struct VaapiH264Decoder {
    info:           NACodecInfoRef,
    vaapi:          Option<VaapiInternals>,
    spses:          Vec<SeqParameterSet>,
    ppses:          Vec<PicParameterSet>,
    frame_refs:     FrameRefs,
    nal_len:        u8,
    out_frm:        NABufferType,
    reorderer:      Reorderer,
    tb_num:         u32,
    tb_den:         u32,
}

fn fill_frame(ifmt: VAImageFormat, pic: &Picture<PictureSync>, frm: &mut NABufferType) -> DecoderResult<()> {
    let mut vbuf = frm.get_vbuf().unwrap();
    let (w, h) = pic.surface_size();
    //let cur_ts = pic.timestamp();

    let img = Image::new(pic, ifmt, w, h, true).expect("get image");

    let iimg = img.image();
    let imgdata: &[u8] = img.as_ref();

    match iimg.format.fourcc().map_err(|_| DecoderError::InvalidData)? {
        VAFourcc::NV12 => {
            let frm = NASimpleVideoFrame::from_video_buf(&mut vbuf).unwrap();
            validate!(iimg.width == (frm.width[0] as u16));
            validate!(iimg.height == (frm.height[0] as u16));

            for (dline, sline) in frm.data[frm.offset[0]..].chunks_mut(frm.stride[0])
                    .zip(imgdata[iimg.offsets[0] as usize..].chunks(iimg.pitches[0] as usize))
                    .take(frm.height[0]) {
                dline[..frm.width[0]].copy_from_slice(&sline[..frm.width[0]]);
            }

            let mut uoff = frm.offset[1];
            let mut voff = frm.offset[2];
            for cline in imgdata[iimg.offsets[1] as usize..].chunks(iimg.pitches[1] as usize).take(frm.height[1]) {
                for (x, pair) in cline.chunks_exact(2).take(frm.width[1]).enumerate() {
                    frm.data[uoff + x] = pair[0];
                    frm.data[voff + x] = pair[1];
                }
                uoff += frm.stride[1];
                voff += frm.stride[2];
            }
        },
        _ => unimplemented!(),
    };
    Ok(())
}

impl Default for VaapiH264Decoder {
    fn default() -> Self {
        Self {
            info:           NACodecInfoRef::default(),
            vaapi:          None,
            spses:          Vec::with_capacity(1),
            ppses:          Vec::with_capacity(4),
            frame_refs:     FrameRefs::new(),
            nal_len:        0,
            out_frm:        NABufferType::None,
            reorderer:      Reorderer::default(),
            tb_num:         0,
            tb_den:         0,
        }
    }
}

impl VaapiH264Decoder {
    pub fn new() -> Self { Self::default() }
    pub fn init(&mut self, info: NACodecInfoRef) -> DecoderResult<()> {
        if let NACodecTypeInfo::Video(vinfo) = info.get_properties() {
            let edata = info.get_extradata().unwrap();
//print!("edata:"); for &el in edata.iter() { print!(" {:02X}", el); } println!();
            let profile;
            let mut nal_buf = Vec::with_capacity(1024);
            if edata.len() > 11 && &edata[0..4] == b"avcC" {
                let mut mr = MemoryReader::new_read(edata.as_slice());
                let mut br = ByteReader::new(&mut mr);

                                          br.read_skip(4)?;
                let version             = br.read_byte()?;
                validate!(version == 1);
                profile                 = br.read_byte()?;
                let _compatibility      = br.read_byte()?;
                let _level              = br.read_byte()?;
                let b                   = br.read_byte()?;
                validate!((b & 0xFC) == 0xFC);
                self.nal_len            = (b & 3) + 1;
                let b                   = br.read_byte()?;
                validate!((b & 0xE0) == 0xE0);
                let num_sps = (b & 0x1F) as usize;
                for _ in 0..num_sps {
                    let len             = br.read_u16be()? as usize;
                    let offset = br.tell() as usize;
                    validate!((br.peek_byte()? & 0x1F) == 7);
                    let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
                                          br.read_skip(len)?;
                    let sps = parse_sps(&nal_buf[1..])?;
                    self.spses.push(sps);
                }
                let num_pps             = br.read_byte()? as usize;
                for _ in 0..num_pps {
                    let len             = br.read_u16be()? as usize;
                    let offset = br.tell() as usize;
                    validate!((br.peek_byte()? & 0x1F) == 8);
                    let _size = unescape_nal(&edata[offset..][..len], &mut nal_buf);
                                          br.read_skip(len)?;
                    let src = &nal_buf;

                    let mut full_size = src.len() * 8;
                    for &byte in src.iter().rev() {
                        if byte == 0 {
                            full_size -= 8;
                        } else {
                            full_size -= (byte.trailing_zeros() + 1) as usize;
                            break;
                        }
                    }
                    validate!(full_size > 0);

                    let pps = parse_pps(&src[1..], &self.spses, full_size - 8)?;
                    let mut found = false;
                    for stored_pps in self.ppses.iter_mut() {
                        if stored_pps.pic_parameter_set_id == pps.pic_parameter_set_id {
                            *stored_pps = pps.clone();
                            found = true;
                            break;
                        }
                    }
                    if !found {
                        self.ppses.push(pps);
                    }
                }
                if br.left() > 0 {
                    match profile {
                        100 | 110 | 122 | 144 => {
                            let b       = br.read_byte()?;
                            validate!((b & 0xFC) == 0xFC);
                            // b & 3 -> chroma format
                            let b       = br.read_byte()?;
                            validate!((b & 0xF8) == 0xF8);
                            // b & 7 -> luma depth minus 8
                            let b       = br.read_byte()?;
                            validate!((b & 0xF8) == 0xF8);
                            // b & 7 -> chroma depth minus 8
                            let num_spsext  = br.read_byte()? as usize;
                            for _ in 0..num_spsext {
                                let len = br.read_u16be()? as usize;
                                // parse spsext
                                          br.read_skip(len)?;
                            }
                        },
                        _ => {},
                    };
                }
            } else {
                return Err(DecoderError::NotImplemented);
            }

            validate!(profile > 0);
            let width  = (vinfo.get_width()  + 15) & !15;
            let height = (vinfo.get_height() + 15) & !15;

            let display = Display::open_silently().expect("open display");

            let num_surfaces = self.spses[0].num_ref_frames + 4 + 64;

            let va_profile = match profile {
                    66 => VAProfile::VAProfileH264ConstrainedBaseline,
                    77 => VAProfile::VAProfileH264Main,
                    88 | 100 | 110 | 122 => VAProfile::VAProfileH264High,
                    _ => return Err(DecoderError::NotImplemented),
                };
            if let Ok(profiles) = display.query_config_profiles() {
                if !profiles.contains(&va_profile) {
println!("Profile {} ({}) not supported", profile, profile_name(va_profile));
                    return Err(DecoderError::NotImplemented);
                }
            } else {
                return Err(DecoderError::Bug);
            }
            if let Ok(points) = display.query_config_entrypoints(va_profile) {
                if !points.contains(&VAEntrypoint::VAEntrypointVLD) {
println!("no decoding support for this profile");
                    return Err(DecoderError::NotImplemented);
                }
            } else {
                return Err(DecoderError::Bug);
            }

            let config = display.create_config(vec![
                    VAConfigAttrib { type_: VAConfigAttribType::VAConfigAttribRTFormat, value: RTFormat::YUV420.into() },
                ], va_profile, VAEntrypoint::VAEntrypointVLD).map_err(|_| {
println!("config creation failed!");
                    DecoderError::Bug
                })?;
            let surfaces = display.create_surfaces(RTFormat::YUV420, None, width as u32, height as u32, Some(UsageHint::Decoder.into()), num_surfaces as u32).map_err(|_| DecoderError::AllocError)?;
            let context = display.create_context(&config, width as i32, height as i32, Some(&surfaces), true).map_err(|_| DecoderError::Bug)?;

            let ref_pics = Vec::new();

            let image_formats = display.query_image_formats().map_err(|_| DecoderError::Bug)?;
            validate!(!image_formats.is_empty());
            let mut ifmt = image_formats[0];
            for fmt in image_formats.iter() {
                if fmt.bits_per_pixel == 12 {
                    ifmt = *fmt;
                    break;
                }
            }

            self.vaapi = Some(VaapiInternals { display, context, ref_pics, surfaces, ifmt });

            let vinfo = NAVideoInfo::new(width, height, false, YUV420_FORMAT);
            self.info = NACodecInfo::new_ref(info.get_name(), NACodecTypeInfo::Video(vinfo), info.get_extradata()).into_ref();
            self.out_frm = alloc_video_buffer(vinfo, 4)?;

            Ok(())
        } else {
            Err(DecoderError::InvalidData)
        }
    }
    fn decode(&mut self, pkt: &NAPacket) -> DecoderResult<()> {
        let src = pkt.get_buffer();
        let vactx = if let Some(ref mut ctx) = self.vaapi { ctx } else { return Err(DecoderError::Bug) };

        let timestamp = pkt.get_dts().unwrap_or_else(|| pkt.get_pts().unwrap_or(0));

        if vactx.surfaces.is_empty() {
panic!("ran out of free surfaces");
//            return Err(DecoderError::AllocError);
        }
        let surface = vactx.surfaces.pop().unwrap();
        let surface_id = surface.id();
        let mut pic = Picture::new(timestamp, vactx.context.clone(), surface);
        let mut is_ref = false;
        let mut is_keyframe = false;

        self.tb_num = pkt.ts.tb_num;
        self.tb_den = pkt.ts.tb_den;

        let mut mr = MemoryReader::new_read(&src);
        let mut br = ByteReader::new(&mut mr);
        let mut frame_type = FrameType::I;
        let mut nal_buf = Vec::with_capacity(1024);
        while br.left() > 0 {
            let size = match self.nal_len {
                    1 => br.read_byte()? as usize,
                    2 => br.read_u16be()? as usize,
                    3 => br.read_u24be()? as usize,
                    4 => br.read_u32be()? as usize,
                    _ => unreachable!(),
                };
            validate!(br.left() >= (size as i64));
            let offset = br.tell() as usize;
            let raw_nal = &src[offset..][..size];
            let _size = unescape_nal(raw_nal, &mut nal_buf);

            let src = &nal_buf;
            validate!((src[0] & 0x80) == 0);
            let nal_ref_idc   = src[0] >> 5;
            let nal_unit_type = src[0] & 0x1F;

            let mut full_size = src.len() * 8;
            for &byte in src.iter().rev() {
                if byte == 0 {
                    full_size -= 8;
                } else {
                    full_size -= (byte.trailing_zeros() + 1) as usize;
                    break;
                }
            }
            validate!(full_size > 0);

            match nal_unit_type {
                 1 | 5 => {
                    let is_idr = nal_unit_type == 5;
                    is_ref |= nal_ref_idc != 0;
                    is_keyframe |= is_idr;
                    let mut br = BitReader::new(&src[..(full_size + 7)/8], BitReaderMode::BE);
                                                        br.skip(8)?;

                    let slice_hdr = parse_slice_header(&mut br, &self.spses, &self.ppses, is_idr, nal_ref_idc)?;
                    match slice_hdr.slice_type {
                        SliceType::P if frame_type != FrameType::B => frame_type = FrameType::P,
                        SliceType::SP if frame_type != FrameType::B => frame_type = FrameType::P,
                        SliceType::B => frame_type = FrameType::B,
                        _ => {},
                    };
                    let mut cur_sps = 0;
                    let mut cur_pps = 0;
                    let mut pps_found = false;
                    for (i, pps) in self.ppses.iter().enumerate() {
                        if pps.pic_parameter_set_id == slice_hdr.pic_parameter_set_id {
                            cur_pps = i;
                            pps_found = true;
                            break;
                        }
                    }
                    validate!(pps_found);
                    let mut sps_found = false;
                    for (i, sps) in self.spses.iter().enumerate() {
                        if sps.seq_parameter_set_id == self.ppses[cur_pps].seq_parameter_set_id {
                            cur_sps = i;
                            sps_found = true;
                            break;
                        }
                    }
                    validate!(sps_found);
                    let sps = &self.spses[cur_sps];
                    let pps = &self.ppses[cur_pps];

                    if slice_hdr.first_mb_in_slice == 0 {
                        let (top_id, bot_id) = self.frame_refs.calc_picture_num(&slice_hdr, is_idr, nal_ref_idc, sps);
                        if is_idr {
                            self.frame_refs.clear_refs();
                            for (pic, _) in vactx.ref_pics.drain(..) {
                                if let Ok(surf) = pic.take_surface() {
                                    vactx.surfaces.push(surf);
                                } else {
                                    panic!("can't take surface");
                                }
                            }
                        }
                        self.frame_refs.select_refs(sps, &slice_hdr, top_id);
                        let mut pic_refs = Vec::with_capacity(NUM_REF_PICS);
                        for pic in self.frame_refs.ref_pics.iter().rev().take(NUM_REF_PICS) {
                            pic_refs.push(pic.make_pic());
                        }
                        if slice_hdr.adaptive_ref_pic_marking_mode {
                            self.frame_refs.apply_adaptive_marking(&slice_hdr.adaptive_ref_pic_marking, slice_hdr.frame_num, 1 << sps.log2_max_frame_num)?;
                        }

                        while pic_refs.len() < NUM_REF_PICS {
                            pic_refs.push(make_dummy_h264_pic());
                        }

                        let mut flags = H264PictureFlags::default();
                        let frame_idx = if let Some(id) = get_long_term_id(is_idr, &slice_hdr) {
                                flags |= H264PictureFlag::LongTermReference;
                                id as u32
                            } else {
                                if nal_ref_idc != 0 {
                                    flags |= H264PictureFlag::ShortTermReference;
                                }
                                u32::from(slice_hdr.frame_num)
                            };
                        let pic_refs: [PictureH264; NUM_REF_PICS] = pic_refs.try_into().unwrap_or_else(|_| panic!("can't convert"));

                        let h264pic = PictureH264::new(surface_id, frame_idx, flags, top_id as i32, bot_id as i32);

                        let seq_fields = H264SeqFields::new(
                                u32::from(sps.chroma_format_idc),
                                u32::from(sps.separate_colour_plane),
                                u32::from(sps.gaps_in_frame_num_value_allowed),
                                u32::from(sps.frame_mbs_only),
                                u32::from(sps.mb_adaptive_frame_field),
                                u32::from(sps.direct_8x8_inference),
                                u32::from(sps.level_idc >= 31),
                                u32::from(sps.log2_max_frame_num) - 4,
                                u32::from(sps.pic_order_cnt_type),
                                u32::from(sps.log2_max_pic_order_cnt_lsb).wrapping_sub(4),
                                u32::from(sps.delta_pic_order_always_zero)
                            );
                        let pic_fields = H264PicFields::new(
                                u32::from(pps.entropy_coding_mode),
                                u32::from(pps.weighted_pred),
                                u32::from(pps.weighted_bipred_idc),
                                u32::from(pps.transform_8x8_mode),
                                u32::from(slice_hdr.field_pic),
                                u32::from(pps.constrained_intra_pred),
                                u32::from(pps.pic_order_present),
                                u32::from(pps.deblocking_filter_control_present),
                                u32::from(pps.redundant_pic_cnt_present),
                                u32::from(nal_ref_idc != 0)
                            );
                        let ppd = PictureParameterBufferH264::new(
                                h264pic,
                                pic_refs,
                                sps.pic_width_in_mbs as u16 - 1,
                                sps.pic_height_in_mbs as u16 - 1,
                                sps.bit_depth_luma - 8,
                                sps.bit_depth_chroma - 8,
                                sps.num_ref_frames as u8,
                                &seq_fields,
                                pps.num_slice_groups as u8 - 1, // should be 0
                                pps.slice_group_map_type, // should be 0
                                0, //pps.slice_group_change_rate as u16 - 1,
                                pps.pic_init_qp as i8 - 26,
                                pps.pic_init_qs as i8 - 26,
                                pps.chroma_qp_index_offset,
                                pps.second_chroma_qp_index_offset,
                                &pic_fields,
                                slice_hdr.frame_num
                            );
                        let pic_param = BufferType::PictureParameter(PictureParameter::H264(ppd));
                        let buf = vactx.context.create_buffer(pic_param).map_err(|_| DecoderError::Bug)?;
                        pic.add_buffer(buf);

                        let mut scaling_list_8x8 = [[0; 64]; 2];
                        scaling_list_8x8[0].copy_from_slice(&pps.scaling_list_8x8[0]);
                        scaling_list_8x8[1].copy_from_slice(&pps.scaling_list_8x8[3]);
                        let iqmatrix = BufferType::IQMatrix(IQMatrix::H264(IQMatrixBufferH264::new(pps.scaling_list_4x4, scaling_list_8x8)));
                        let buf = vactx.context.create_buffer(iqmatrix).map_err(|_| DecoderError::Bug)?;
                        pic.add_buffer(buf);

                        let cpic = PictureInfo {
                                id: slice_hdr.frame_num,
                                full_id: top_id,
                                surface_id,
                                top_id, bot_id,
                                //pic_type: slice_hdr.slice_type.to_frame_type(),
                                is_ref,
                                is_idr,
                                long_term: get_long_term_id(is_idr, &slice_hdr),
                            };
                        if cpic.is_ref {
                            self.frame_refs.add_short_term(cpic.clone(), sps.num_ref_frames);
                        }
                        if let Some(lt_idx) = cpic.long_term {
                            self.frame_refs.add_long_term(lt_idx, cpic);
                        }
                    }

                    let mut luma_weight_l0 = [0i16; 32];
                    let mut luma_offset_l0 = [0i16; 32];
                    let mut chroma_weight_l0 = [[0i16; 2]; 32];
                    let mut chroma_offset_l0 = [[0i16; 2]; 32];
                    let mut luma_weight_l1 = [0i16; 32];
                    let mut luma_offset_l1 = [0i16; 32];
                    let mut chroma_weight_l1 = [[0i16; 2]; 32];
                    let mut chroma_offset_l1 = [[0i16; 2]; 32];
                    let mut luma_weighted_l0 = false;
                    let mut chroma_weighted_l0 = false;
                    let mut luma_weighted_l1 = false;
                    let mut chroma_weighted_l1 = false;
                    let mut luma_log2_weight_denom = slice_hdr.luma_log2_weight_denom;
                    let mut chroma_log2_weight_denom = slice_hdr.chroma_log2_weight_denom;

                    if (pps.weighted_pred && matches!(slice_hdr.slice_type, SliceType::P | SliceType::B)) || (pps.weighted_bipred_idc == 1 && slice_hdr.slice_type == SliceType::B) {
                        luma_weighted_l0 = true;
                        chroma_weighted_l0 = false;
                        for (i, winfo) in slice_hdr.weights_l0.iter().enumerate().take(slice_hdr.num_ref_idx_l0_active) {
                            if winfo.luma_weighted {
                                luma_weight_l0[i] = winfo.luma_weight.into();
                                luma_offset_l0[i] = winfo.luma_offset.into();
                            } else {
                                luma_weight_l0[i] = 1 << slice_hdr.luma_log2_weight_denom;
                            }
                            if winfo.chroma_weighted {
                                chroma_weight_l0[i][0] = winfo.chroma_weight[0].into();
                                chroma_weight_l0[i][1] = winfo.chroma_weight[1].into();
                                chroma_offset_l0[i][0] = winfo.chroma_offset[0].into();
                                chroma_offset_l0[i][1] = winfo.chroma_offset[1].into();
                            } else {
                                chroma_weight_l0[i][0] = 1 << slice_hdr.chroma_log2_weight_denom;
                                chroma_weight_l0[i][1] = 1 << slice_hdr.chroma_log2_weight_denom;
                                chroma_offset_l0[i][0] = 0;
                                chroma_offset_l0[i][1] = 0;
                            }
                            chroma_weighted_l0 |= winfo.chroma_weighted;
                        }
                    }
                    if pps.weighted_bipred_idc == 1 && slice_hdr.slice_type == SliceType::B {
                        luma_weighted_l1 = true;
                        chroma_weighted_l1 = sps.chroma_format_idc != 0;
                        for (i, winfo) in slice_hdr.weights_l1.iter().enumerate().take(slice_hdr.num_ref_idx_l1_active) {
                            if winfo.luma_weighted {
                                luma_weight_l1[i] = winfo.luma_weight.into();
                                luma_offset_l1[i] = winfo.luma_offset.into();
                            } else {
                                luma_weight_l1[i] = 1 << slice_hdr.luma_log2_weight_denom;
                            }
                            if chroma_weighted_l1 && winfo.chroma_weighted {
                                chroma_weight_l1[i][0] = winfo.chroma_weight[0].into();
                                chroma_weight_l1[i][1] = winfo.chroma_weight[1].into();
                                chroma_offset_l1[i][0] = winfo.chroma_offset[0].into();
                                chroma_offset_l1[i][1] = winfo.chroma_offset[1].into();
                            } else {
                                chroma_weight_l1[i][0] = 1 << slice_hdr.chroma_log2_weight_denom;
                                chroma_weight_l1[i][1] = 1 << slice_hdr.chroma_log2_weight_denom;
                                chroma_offset_l1[i][0] = 0;
                                chroma_offset_l1[i][1] = 0;
                            }
                        }
                    }
                    if pps.weighted_bipred_idc == 2 && slice_hdr.slice_type == SliceType::B {
                        let num_l0 = slice_hdr.num_ref_idx_l0_active;
                        let num_l1 = slice_hdr.num_ref_idx_l1_active;
                        if num_l0 != 1 || num_l1 != 1 { //xxx: also exclude symmetric case
                            luma_weighted_l0 = false;
                            luma_weighted_l1 = false;
                            chroma_weighted_l0 = false;
                            chroma_weighted_l1 = false;
                            luma_log2_weight_denom = 5;
                            chroma_log2_weight_denom = 5;

                            for w in luma_weight_l0.iter_mut() {
                                *w = 32;
                            }
                            for w in luma_weight_l1.iter_mut() {
                                *w = 32;
                            }
                            for w in chroma_weight_l0.iter_mut() {
                                *w = [32; 2];
                            }
                            for w in chroma_weight_l1.iter_mut() {
                                *w = [32; 2];
                            }
                        }
                    }

                    let ref_pic_list_0 = map_ref_list(&self.frame_refs.cur_refs.ref_list0);
                    let ref_pic_list_1 = map_ref_list(&self.frame_refs.cur_refs.ref_list1);

                    let slice_param = SliceParameterBufferH264::new(
                            raw_nal.len() as u32,
                            0, // no offset
                            VASliceDataFlag::All,
                            br.tell() as u16,
                            slice_hdr.first_mb_in_slice as u16,
                            match slice_hdr.slice_type {
                                SliceType::I => 2,
                                SliceType::P => 0,
                                SliceType::B => 1,
                                SliceType::SI => 4,
                                SliceType::SP => 3,
                            },
                            slice_hdr.direct_spatial_mv_pred as u8,
                            (slice_hdr.num_ref_idx_l0_active as u8).saturating_sub(1),
                            (slice_hdr.num_ref_idx_l1_active as u8).saturating_sub(1),
                            slice_hdr.cabac_init_idc,
                            slice_hdr.slice_qp_delta as i8,
                            slice_hdr.disable_deblocking_filter_idc,
                            slice_hdr.slice_alpha_c0_offset / 2,
                            slice_hdr.slice_beta_offset / 2,
                            ref_pic_list_0,
                            ref_pic_list_1,
                            luma_log2_weight_denom,
                            chroma_log2_weight_denom,
                            luma_weighted_l0 as u8, luma_weight_l0, luma_offset_l0,
                            chroma_weighted_l0 as u8, chroma_weight_l0, chroma_offset_l0,
                            luma_weighted_l1 as u8, luma_weight_l1, luma_offset_l1,
                            chroma_weighted_l1 as u8, chroma_weight_l1, chroma_offset_l1,
                        );
                    let slc_param = BufferType::SliceParameter(SliceParameter::H264(slice_param));
                    let buf = vactx.context.create_buffer(slc_param).map_err(|_| DecoderError::Bug)?;
                    pic.add_buffer(buf);

                    let slc_data = BufferType::SliceData(raw_nal.to_vec());
                    let buf = vactx.context.create_buffer(slc_data).map_err(|_| DecoderError::Bug)?;
                    pic.add_buffer(buf);
                },
                 2 => { // slice data partition A
                    //slice header
                    //slice id = read_ue()
                    //cat 2 slice data (all but MB layer residual)
                    return Err(DecoderError::NotImplemented);
                },
                 3 => { // slice data partition B
                    //slice id = read_ue()
                    //if pps.redundant_pic_cnt_present { redundant_pic_cnt = read_ue() }
                    //cat 3 slice data (MB layer residual)
                    return Err(DecoderError::NotImplemented);
                },
                 4 => { // slice data partition C
                    //slice id = read_ue()
                    //if pps.redundant_pic_cnt_present { redundant_pic_cnt = read_ue() }
                    //cat 4 slice data (MB layer residual)
                    return Err(DecoderError::NotImplemented);
                },
                 6 => {}, //SEI
                 7 => {
                    let sps = parse_sps(&src[1..])?;
                    self.spses.push(sps);
                },
                 8 => {
                    validate!(full_size >= 8 + 16);
                    let pps = parse_pps(&src[1..], &self.spses, full_size - 8)?;
                    let mut found = false;
                    for stored_pps in self.ppses.iter_mut() {
                        if stored_pps.pic_parameter_set_id == pps.pic_parameter_set_id {
                            *stored_pps = pps.clone();
                            found = true;
                            break;
                        }
                    }
                    if !found {
                        self.ppses.push(pps);
                    }
                },
                 9 => { // access unit delimiter
                },
                10 => {}, //end of sequence
                11 => {}, //end of stream
                12 => {}, //filler
                _  => {},
            };

            br.read_skip(size)?;
        }

        let bpic = pic.begin().expect("begin");
        let rpic = bpic.render().expect("render");
        let epic = rpic.end().expect("end");

        self.reorderer.add_frame(WaitingFrame {
                pic:    epic,
                is_idr: is_keyframe,
                is_ref,
                ftype:  frame_type,
                ts:     timestamp,
            });

        let mut idx = 0;
        while idx < vactx.ref_pics.len() {
            let cur_surf_id = vactx.ref_pics[idx].1;
            if self.frame_refs.ref_pics.iter().any(|fref| fref.surface_id == cur_surf_id) {
                idx += 1;
            } else {
                let (pic, _) = vactx.ref_pics.remove(idx);
                if let Ok(surf) = pic.take_surface() {
                    vactx.surfaces.push(surf);
                } else {
                    panic!("can't take surface");
                }
            }
        }

        Ok(())
    }
    fn get_frame(&mut self) -> Option<NAFrameRef> {
        if let Some(ref mut vactx) = self.vaapi {
            if let Some(frm) = self.reorderer.get_frame() {
                let ts = frm.ts;
                let is_idr = frm.is_idr;
                let is_ref = frm.is_ref;
                let ftype = frm.ftype;
                if let Ok(pic) = frm.pic.sync() {
                    let _ = fill_frame(vactx.ifmt, &pic, &mut self.out_frm);

                    if !is_ref {
                        if let Ok(surf) = pic.take_surface() {
                            vactx.surfaces.push(surf);
                        } else {
                            panic!("can't take surface");
                        }
                    } else {
                        let id = pic.surface_id();
                        vactx.ref_pics.push((pic, id));
                    }

                    let ts = NATimeInfo::new(None, Some(ts), None, self.tb_num, self.tb_den);
                    Some(NAFrame::new(ts, ftype, is_idr, self.info.clone(), self.out_frm.clone()).into_ref())
                } else {
                    panic!("can't sync");
                }
            } else {
                None
            }
        } else {
            None
        }
    }
    fn get_last_frames(&mut self) -> Option<NAFrameRef> {
        if let Some(ref mut vactx) = self.vaapi {
            if let Some(frm) = self.reorderer.frames.pop_front() {
                let ts = frm.ts;
                let is_idr = frm.is_idr;
                let is_ref = frm.is_ref;
                let ftype = frm.ftype;
                if let Ok(pic) = frm.pic.sync() {
                    let _ = fill_frame(vactx.ifmt, &pic, &mut self.out_frm);

                    if !is_ref {
                        if let Ok(surf) = pic.take_surface() {
                            vactx.surfaces.push(surf);
                        } else {
                            panic!("can't take surface");
                        }
                    } else {
                        let id = pic.surface_id();
                        vactx.ref_pics.push((pic, id));
                    }

                    let ts = NATimeInfo::new(None, Some(ts), None, self.tb_num, self.tb_den);
                    Some(NAFrame::new(ts, ftype, is_idr, self.info.clone(), self.out_frm.clone()).into_ref())
                } else {
                    panic!("can't sync");
                }
            } else {
                None
            }
        } else {
            None
        }
    }
    fn flush(&mut self) {
        self.frame_refs.clear_refs();
        if let Some(ref mut vactx) = self.vaapi {
            for frm in self.reorderer.frames.drain(..) {
                if let Ok(pic) = frm.pic.sync() {
                    if let Ok(surf) = pic.take_surface() {
                        vactx.surfaces.push(surf);
                    } else {
                        panic!("can't take surface");
                    }
                } else {
                    panic!("can't sync");
                }
            }
            self.reorderer.flush();
            for (pic, _) in vactx.ref_pics.drain(..) {
                if let Ok(surf) = pic.take_surface() {
                    vactx.surfaces.push(surf);
                } else {
                    panic!("can't take surface");
                }
            }
        }
    }
}

impl NAOptionHandler for VaapiH264Decoder {
    fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
    fn set_options(&mut self, _options: &[NAOption]) {}
    fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
}

use std::thread::*;
use std::sync::mpsc::*;

enum DecMessage {
    Init(NACodecInfoRef),
    Decode(NAPacket),
    Flush,
    GetFrame,
    GetLastFrames,
    End
}

enum DecResponse {
    Ok,
    Nothing,
    Err(DecoderError),
    Frame(NAFrameRef),
}

pub trait HWDecoder {
    fn init(&mut self, info: NACodecInfoRef) -> DecoderResult<()>;
    fn queue_pkt(&mut self, pkt: &NAPacket) -> DecoderResult<()>;
    fn get_frame(&mut self) -> Option<NAFrameRef>;
    fn get_last_frames(&mut self) -> Option<NAFrameRef>;
    fn flush(&mut self);
}

pub struct HWWrapper {
    handle:     Option<JoinHandle<DecoderResult<()>>>,
    send:       SyncSender<DecMessage>,
    recv:       Receiver<DecResponse>,
}

#[allow(clippy::new_without_default)]
impl HWWrapper {
    pub fn new() -> Self {
        let (in_send, in_recv) = sync_channel(1);
        let (out_send, out_recv) = sync_channel(1);
        let handle = std::thread::spawn(move || {
                let receiver = in_recv;
                let sender = out_send;
                let mut dec = VaapiH264Decoder::new();
                while let Ok(msg) = receiver.recv() {
                    match msg {
                        DecMessage::Init(info) => {
                            let msg = if let Err(err) = dec.init(info) {
                                    DecResponse::Err(err)
                                } else {
                                    DecResponse::Ok
                                };
                            sender.send(msg).map_err(|_| DecoderError::Bug)?;
                        },
                        DecMessage::Decode(pkt) => {
                            let msg = match dec.decode(&pkt) {
                                    Ok(()) => DecResponse::Ok,
                                    Err(err) => DecResponse::Err(err),
                                };
                            sender.send(msg).map_err(|_| DecoderError::Bug)?;
                        },
                        DecMessage::GetFrame => {
                            let msg = match dec.get_frame() {
                                    Some(frm) => DecResponse::Frame(frm),
                                    None => DecResponse::Nothing,
                                };
                            sender.send(msg).map_err(|_| DecoderError::Bug)?;
                        },
                        DecMessage::GetLastFrames => {
                            let msg = match dec.get_last_frames() {
                                    Some(frm) => DecResponse::Frame(frm),
                                    None => DecResponse::Nothing,
                                };
                            sender.send(msg).map_err(|_| DecoderError::Bug)?;
                        },
                        DecMessage::Flush => dec.flush(),
                        DecMessage::End => return Ok(()),
                    };
                }
                Err(DecoderError::Bug)
            });

        Self {
            handle:     Some(handle),
            send:       in_send,
            recv:       out_recv,
        }
    }
}

impl HWDecoder for HWWrapper {
    fn init(&mut self, info: NACodecInfoRef) -> DecoderResult<()> {
        if self.send.send(DecMessage::Init(info)).is_ok() {
            match self.recv.recv() {
                Ok(DecResponse::Ok) => Ok(()),
                Ok(DecResponse::Err(err)) => Err(err),
                Err(_) => Err(DecoderError::Bug),
                _ => unreachable!(),
            }
        } else {
            Err(DecoderError::Bug)
        }
    }
    fn queue_pkt(&mut self, pkt: &NAPacket) -> DecoderResult<()> {
        let pkt2 = NAPacket::new_from_refbuf(pkt.get_stream(), pkt.ts, pkt.keyframe, pkt.get_buffer());
        if self.send.send(DecMessage::Decode(pkt2)).is_ok() {
            match self.recv.recv() {
                Ok(DecResponse::Ok) => Ok(()),
                Ok(DecResponse::Err(err)) => Err(err),
                Err(_) => Err(DecoderError::Bug),
                _ => unreachable!(),
            }
        } else {
            Err(DecoderError::Bug)
        }
    }
    fn get_frame(&mut self) -> Option<NAFrameRef> {
        if self.send.send(DecMessage::GetFrame).is_ok() {
            match self.recv.recv() {
                Ok(DecResponse::Frame(frm)) => Some(frm),
                Ok(DecResponse::Nothing) => None,
                Err(_) => None,
                _ => unreachable!(),
            }
        } else {
            None
        }
    }
    fn get_last_frames(&mut self) -> Option<NAFrameRef> {
        if self.send.send(DecMessage::GetLastFrames).is_ok() {
            match self.recv.recv() {
                Ok(DecResponse::Frame(frm)) => Some(frm),
                Ok(DecResponse::Nothing) => None,
                Err(_) => None,
                _ => unreachable!(),
            }
        } else {
            None
        }
    }
    fn flush(&mut self) {
        let _ = self.send.send(DecMessage::Flush);
    }
}

impl Drop for HWWrapper {
    fn drop(&mut self) {
        if self.send.send(DecMessage::End).is_ok() {
            let mut handle = None;
            std::mem::swap(&mut handle, &mut self.handle);
            if let Some(hdl) = handle {
                let _ = hdl.join();
            }
        }
    }
}

impl NAOptionHandler for HWWrapper {
    fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
    fn set_options(&mut self, _options: &[NAOption]) {}
    fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
}

pub fn new_h264_hwdec() -> Box<dyn HWDecoder + Send> {
    Box::new(HWWrapper::new())
}

#[cfg(test)]
mod test {
    use nihav_core::codecs::*;
    use nihav_core::io::byteio::*;
    use nihav_core::demuxers::{RegisteredDemuxers, create_demuxer};
    use nihav_commonfmt::generic_register_all_demuxers;
    use super::VaapiH264Decoder;
    use std::io::prelude::*;

    fn decode_h264(name: &str, dname: &str, dmx_reg: &RegisteredDemuxers, opfx: &str) -> DecoderResult<()> {
        let dmx_f = dmx_reg.find_demuxer(dname).expect("demuxer exists");
        let file = std::fs::File::open(name).expect("file exists");
        let mut fr = FileReader::new_read(file);
        let mut br = ByteReader::new(&mut fr);
        let mut dmx = create_demuxer(dmx_f, &mut br).expect("create demuxer");

        let mut vstream_id = 0;
        let mut dec = VaapiH264Decoder::new();
        for stream in dmx.get_streams() {
            if stream.get_media_type() == StreamType::Video {
                dec.init(stream.get_info()).expect("inited");
                vstream_id = stream.get_id();
                break;
            }
        }

        let mut frameno = 0;
        while let Ok(pkt) = dmx.get_frame() {
            if pkt.get_stream().get_id() != vstream_id {
                continue;
            }
            dec.decode(&pkt).expect("decoded");
            let frm = dec.get_last_frames().expect("get frame");
            let timestamp = frm.get_dts().unwrap_or_else(|| frm.get_pts().unwrap_or(0));

            let pic = frm.get_buffer().get_vbuf().expect("got picture");

            let nname = format!("assets/test_out/{}{:06}_{}.pgm", opfx, timestamp, frameno);
            frameno += 1;
            let mut file = std::fs::File::create(&nname).expect("create file");
            let (w, h) = pic.get_dimensions(0);
            file.write_all(format!("P5\n{} {}\n255\n", w, h * 3 / 2).as_bytes()).expect("header written");
            let data = pic.get_data();
            for yline in data.chunks(pic.get_stride(0)).take(h) {
                file.write_all(&yline[..w]).expect("Y line written");
            }
            for (uline, vline) in data[pic.get_offset(1)..].chunks(pic.get_stride(1))
                    .zip(data[pic.get_offset(2)..].chunks(pic.get_stride(2))).take(h / 2) {
                file.write_all(&uline[..w / 2]).expect("U line written");
                file.write_all(&vline[..w / 2]).expect("V line written");
            }
        }
        Ok(())
    }


    // samples if not specified otherwise come from H.264 conformance suite

    #[test]
    fn test_h264_simple() {
        let mut dmx_reg = RegisteredDemuxers::new();
        generic_register_all_demuxers(&mut dmx_reg);

        decode_h264("assets/ITU/DimpledSpanishCuckoo-mobile.mp4", "mov", &dmx_reg, "hw").unwrap();
    }
}