//! VAAPI H.264 hardware encoding.
//!
//! Two-phase initialization:
//! - Phase 1 (`VaapiEncoder::new`): store config, no ffmpeg calls.
//! - Phase 2 (first call to `encode`): build filter graph using `hw_frames_ctx`
//!   from the first DRM_PRIME frame to wire up hwmap, then open the codec.
//!
//! The split exists because the filter graph needs the DRM hardware device
//! context that kmsgrab attaches to each frame — that context isn't available
//! until the first frame arrives, so the graph can't be validated at
//! construction time.

use anyhow::{Context, Result};
use tracing::info;

/// Encoding configuration.
pub struct EncoderConfig {
    pub width: u32,
    pub height: u32,
    pub fps: u32,
    /// Quantization parameter (lower = higher quality). Default: 20
    pub qp: u32,
    /// Keyframe interval in frames. Default: 30 (1 keyframe/sec at 30fps)
    pub gop_size: u32,
}

impl Default for EncoderConfig {
    fn default() -> Self {
        Self {
            width: 1920,
            height: 1080,
            fps: 30,
            qp: 20,
            gop_size: 30,
        }
    }
}

/// VAAPI H.264 encoder with lazy filter graph initialization.
pub struct VaapiEncoder {
    config: EncoderConfig,
    inner: Option<EncoderInner>,
}

struct EncoderInner {
    encoder: ffmpeg::encoder::Video,
    filter_graph: ffmpeg::filter::Graph,
    frame_count: u64,
    time_base: ffmpeg::Rational,
}

impl VaapiEncoder {
    /// Create an encoder from config. No ffmpeg resources are allocated yet.
    pub fn new(config: EncoderConfig) -> Self {
        Self { config, inner: None }
    }

    /// Encode a frame. On the first call, initializes the filter graph and
    /// codec using the frame's `hw_frames_ctx`. Returns 0 or more packets.
    pub fn encode(&mut self, frame: &ffmpeg::frame::Video) -> Result<Vec<EncodedPacket>> {
        if self.inner.is_none() {
            self.inner = Some(EncoderInner::open(&self.config, frame).context("encoder init")?);
        }
        self.inner.as_mut().unwrap().encode(frame)
    }

    /// Flush remaining packets out of the encoder.
    pub fn flush(&mut self) -> Result<Vec<EncodedPacket>> {
        match self.inner.as_mut() {
            Some(inner) => inner.flush(),
            None => Ok(vec![]),
        }
    }
}

impl EncoderInner {
    fn open(config: &EncoderConfig, first_frame: &ffmpeg::frame::Video) -> Result<Self> {
        let time_base = ffmpeg::Rational::new(1, config.fps as i32);
        let filter_graph = Self::build_filter_graph(config, first_frame)?;

        let codec = ffmpeg::encoder::find_by_name("h264_vaapi")
            .context("h264_vaapi encoder not found")?;

        let mut encoder_ctx = ffmpeg::codec::context::Context::new_with_codec(codec)
            .encoder()
            .video()?;

        encoder_ctx.set_width(config.width);
        encoder_ctx.set_height(config.height);
        encoder_ctx.set_time_base(time_base);
        encoder_ctx.set_frame_rate(Some(ffmpeg::Rational::new(config.fps as i32, 1)));
        encoder_ctx.set_gop(config.gop_size);
        encoder_ctx.set_max_b_frames(0);
        encoder_ctx.set_format(ffmpeg::format::Pixel::VAAPI);

        let mut opts = ffmpeg::Dictionary::new();
        opts.set("qp", &config.qp.to_string());

        let encoder = encoder_ctx
            .open_with(opts)
            .context("open h264_vaapi encoder")?;

        info!(
            "VAAPI encoder opened: {}x{} @ {}fps, qp={}, gop={}",
            config.width, config.height, config.fps, config.qp, config.gop_size,
        );

        Ok(Self { encoder, filter_graph, frame_count: 0, time_base })
    }

    /// Build the filter graph: DRM_PRIME → hwmap(vaapi) → scale_vaapi → buffersink.
    ///
    /// Mirrors the ffmpeg CLI approach:
    ///   -init_hw_device drm=drm:/dev/dri/card0
    ///   -init_hw_device vaapi=va@drm
    ///   -filter_hw_device va
    ///
    /// We create both devices explicitly and set the VAAPI device on the graph,
    /// then attach the frame's hw_frames_ctx to the buffersrc so hwmap can map
    /// DRM_PRIME frames to VAAPI surfaces.
    fn build_filter_graph(
        config: &EncoderConfig,
        first_frame: &ffmpeg::frame::Video,
    ) -> Result<ffmpeg::filter::Graph> {
        let input_width = first_frame.width();
        let input_height = first_frame.height();

        // Create DRM device, then derive VAAPI from it (like -init_hw_device vaapi=va@drm)
        let (drm_device, vaapi_device) = unsafe {
            let mut drm_ref: *mut ffmpeg::ffi::AVBufferRef = std::ptr::null_mut();
            let dev_path = std::ffi::CString::new("/dev/dri/card0").unwrap();
            let ret = ffmpeg::ffi::av_hwdevice_ctx_create(
                &mut drm_ref,
                ffmpeg::ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_DRM,
                dev_path.as_ptr(),
                std::ptr::null_mut(),
                0,
            );
            anyhow::ensure!(ret >= 0, "av_hwdevice_ctx_create(drm): {ret}");
            info!("DRM device created");

            let mut vaapi_ref: *mut ffmpeg::ffi::AVBufferRef = std::ptr::null_mut();
            let ret = ffmpeg::ffi::av_hwdevice_ctx_create_derived(
                &mut vaapi_ref,
                ffmpeg::ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
                drm_ref,
                0,
            );
            anyhow::ensure!(ret >= 0, "av_hwdevice_ctx_create_derived(vaapi): {ret}");
            info!("VAAPI device derived from DRM");

            (drm_ref, vaapi_ref)
        };

        let mut graph = ffmpeg::filter::Graph::new();

        let args = format!(
            "video_size={}x{}:pix_fmt={}:time_base=1/{}:pixel_aspect=1/1",
            input_width,
            input_height,
            ffmpeg::format::Pixel::DRM_PRIME as i32,
            config.fps,
        );

        graph
            .add(&ffmpeg::filter::find("buffer").unwrap(), "in", &args)
            .context("add buffersrc")?;
        graph
            .add(&ffmpeg::filter::find("buffersink").unwrap(), "out", "")
            .context("add buffersink")?;

        let filter_spec = format!(
            "hwmap=derive_device=vaapi,scale_vaapi=w={}:h={}:format=nv12",
            config.width, config.height,
        );

        graph.output("in", 0)?.input("out", 0)?.parse(&filter_spec)?;

        // Set the VAAPI device on filter contexts that need it
        // (equivalent to -filter_hw_device va in the CLI)
        unsafe {
            let graph_ptr = graph.as_mut_ptr();
            for i in 0..(*graph_ptr).nb_filters {
                let fctx = *(*graph_ptr).filters.add(i as usize);
                (*fctx).hw_device_ctx = ffmpeg::ffi::av_buffer_ref(vaapi_device);
            }
        }

        // Attach hw_frames_ctx from the first DRM_PRIME frame to the buffersrc
        unsafe {
            let hw_frames_ctx = (*first_frame.as_ptr()).hw_frames_ctx;
            anyhow::ensure!(
                !hw_frames_ctx.is_null(),
                "first KMS frame has no hw_frames_ctx"
            );

            let mut buffersrc = graph.get("in").unwrap();
            let par = ffmpeg::ffi::av_buffersrc_parameters_alloc();
            anyhow::ensure!(!par.is_null(), "av_buffersrc_parameters_alloc OOM");
            (*par).hw_frames_ctx = ffmpeg::ffi::av_buffer_ref(hw_frames_ctx);
            let ret = ffmpeg::ffi::av_buffersrc_parameters_set(buffersrc.as_mut_ptr(), par);
            ffmpeg::ffi::av_free(par as *mut _);
            anyhow::ensure!(ret >= 0, "av_buffersrc_parameters_set: {ret}");
        }

        graph.validate().context("validate filter graph")?;
        info!("Filter graph ready: {filter_spec}");

        // Keep device refs alive — they're ref-counted by the graph now,
        // but we drop our refs here.
        unsafe {
            ffmpeg::ffi::av_buffer_unref(&mut (drm_device as *mut _));
            ffmpeg::ffi::av_buffer_unref(&mut (vaapi_device as *mut _));
        }

        Ok(graph)
    }

    fn encode(&mut self, frame: &ffmpeg::frame::Video) -> Result<Vec<EncodedPacket>> {
        self.filter_graph
            .get("in")
            .unwrap()
            .source()
            .add(frame)
            .context("buffersrc add")?;

        let mut packets = Vec::new();
        let mut filtered = ffmpeg::frame::Video::empty();

        while self
            .filter_graph
            .get("out")
            .unwrap()
            .sink()
            .frame(&mut filtered)
            .is_ok()
        {
            filtered.set_pts(Some(self.frame_count as i64));
            self.frame_count += 1;

            self.encoder.send_frame(&filtered)?;

            let mut encoded = ffmpeg::Packet::empty();
            while self.encoder.receive_packet(&mut encoded).is_ok() {
                packets.push(EncodedPacket {
                    data: encoded.data().unwrap_or(&[]).to_vec(),
                    pts: encoded.pts().unwrap_or(0),
                    dts: encoded.dts().unwrap_or(0),
                    keyframe: encoded.is_key(),
                    time_base_num: self.time_base.numerator() as u32,
                    time_base_den: self.time_base.denominator() as u32,
                });
            }
        }

        Ok(packets)
    }

    fn flush(&mut self) -> Result<Vec<EncodedPacket>> {
        self.encoder.send_eof()?;
        let mut packets = Vec::new();
        let mut encoded = ffmpeg::Packet::empty();
        while self.encoder.receive_packet(&mut encoded).is_ok() {
            packets.push(EncodedPacket {
                data: encoded.data().unwrap_or(&[]).to_vec(),
                pts: encoded.pts().unwrap_or(0),
                dts: encoded.dts().unwrap_or(0),
                keyframe: encoded.is_key(),
                time_base_num: self.time_base.numerator() as u32,
                time_base_den: self.time_base.denominator() as u32,
            });
        }
        Ok(packets)
    }
}

/// An encoded video packet ready for transport.
pub struct EncodedPacket {
    pub data: Vec<u8>,
    pub pts: i64,
    pub dts: i64,
    pub keyframe: bool,
    pub time_base_num: u32,
    pub time_base_den: u32,
}