//! ext-image-copy-capture-v1 and ext-image-capture-source-v1 protocol implementation.
//!
//! This provides the new standardized screencopy protocol that replaces wlr-screencopy.
//! Clients like `grim` and screen recorders will use this when available.

use std::time::UNIX_EPOCH;

use smithay::reexports::wayland_protocols::ext::image_capture_source::v1::server::{
    ext_image_capture_source_v1::{self, ExtImageCaptureSourceV1},
    ext_output_image_capture_source_manager_v1::{
        self, ExtOutputImageCaptureSourceManagerV1,
    },
};
use smithay::reexports::wayland_protocols::ext::image_copy_capture::v1::server::{
    ext_image_copy_capture_cursor_session_v1::{self, ExtImageCopyCaptureCursorSessionV1},
    ext_image_copy_capture_frame_v1::{self, ExtImageCopyCaptureFrameV1},
    ext_image_copy_capture_manager_v1::{self, ExtImageCopyCaptureManagerV1},
    ext_image_copy_capture_session_v1::{self, ExtImageCopyCaptureSessionV1},
};
use smithay::reexports::wayland_server::protocol::wl_buffer::WlBuffer;
use smithay::reexports::wayland_server::{
    Client, DataInit, Dispatch, DisplayHandle, GlobalDispatch, New,
};
use tracing::debug;

const MANAGER_VERSION: u32 = 1;
const SOURCE_MANAGER_VERSION: u32 = 1;

/// Per-source user data.
pub struct CaptureSourceData;

/// Per-session user data.
pub struct CaptureSessionData;

/// Per-cursor-session user data.
pub struct CaptureCursorSessionData;

/// Per-frame user data — stores the attached buffer for pixel copy.
pub struct CaptureFrameData {
    pub buffer: Option<WlBuffer>,
}

/// State for the ext-image-copy-capture protocol.
pub struct ImageCopyCaptureState;

/// A pending image-copy-capture frame queued for the next render.
pub struct ImageCopyCapture {
    pub frame: ExtImageCopyCaptureFrameV1,
    pub buffer: WlBuffer,
    pub submitted: bool,
}

impl Drop for ImageCopyCapture {
    fn drop(&mut self) {
        if !self.submitted {
            self.frame.failed(ext_image_copy_capture_frame_v1::FailureReason::Unknown);
        }
    }
}

impl ImageCopyCapture {
    pub fn buffer(&self) -> &WlBuffer {
        &self.buffer
    }

    /// Signal the frame as ready after pixel data has been copied.
    pub fn submit(mut self) {
        let now = UNIX_EPOCH.elapsed().unwrap();
        let secs = now.as_secs();
        self.frame.presentation_time(
            (secs >> 32) as u32,
            secs as u32,
            now.subsec_nanos(),
        );
        self.frame.ready();
        self.submitted = true;
    }

    /// Signal capture failure.
    pub fn fail(mut self) {
        self.frame.failed(ext_image_copy_capture_frame_v1::FailureReason::Unknown);
        self.submitted = true;
    }
}

pub trait ImageCopyCaptureHandler {
    fn image_copy_frame(&mut self, capture: ImageCopyCapture);
}

impl ImageCopyCaptureState {
    pub fn new<D>(display: &DisplayHandle) -> Self
    where
        D: GlobalDispatch<ExtImageCopyCaptureManagerV1, ()>
            + Dispatch<ExtImageCopyCaptureManagerV1, ()>
            + Dispatch<ExtImageCopyCaptureSessionV1, CaptureSessionData>
            + Dispatch<ExtImageCopyCaptureCursorSessionV1, CaptureCursorSessionData>
            + Dispatch<ExtImageCopyCaptureFrameV1, CaptureFrameData>
            + GlobalDispatch<ExtOutputImageCaptureSourceManagerV1, ()>
            + Dispatch<ExtOutputImageCaptureSourceManagerV1, ()>
            + Dispatch<ExtImageCaptureSourceV1, CaptureSourceData>
            + ImageCopyCaptureHandler
            + 'static,
    {
        display.create_global::<D, ExtImageCopyCaptureManagerV1, _>(MANAGER_VERSION, ());
        display.create_global::<D, ExtOutputImageCaptureSourceManagerV1, _>(
            SOURCE_MANAGER_VERSION,
            (),
        );
        Self
    }
}

// --- ExtOutputImageCaptureSourceManagerV1 (global) ---

impl<D> GlobalDispatch<ExtOutputImageCaptureSourceManagerV1, (), D> for ImageCopyCaptureState
where
    D: GlobalDispatch<ExtOutputImageCaptureSourceManagerV1, ()>
        + Dispatch<ExtOutputImageCaptureSourceManagerV1, ()>
        + Dispatch<ExtImageCaptureSourceV1, CaptureSourceData>
        + 'static,
{
    fn bind(
        _state: &mut D,
        _display: &DisplayHandle,
        _client: &Client,
        resource: New<ExtOutputImageCaptureSourceManagerV1>,
        _data: &(),
        data_init: &mut DataInit<'_, D>,
    ) {
        data_init.init(resource, ());
        debug!("ext-output-image-capture-source-manager bound");
    }
}

impl<D> Dispatch<ExtOutputImageCaptureSourceManagerV1, (), D> for ImageCopyCaptureState
where
    D: Dispatch<ExtOutputImageCaptureSourceManagerV1, ()>
        + Dispatch<ExtImageCaptureSourceV1, CaptureSourceData>
        + 'static,
{
    fn request(
        _state: &mut D,
        _client: &Client,
        _resource: &ExtOutputImageCaptureSourceManagerV1,
        request: ext_output_image_capture_source_manager_v1::Request,
        _data: &(),
        _display: &DisplayHandle,
        data_init: &mut DataInit<'_, D>,
    ) {
        match request {
            ext_output_image_capture_source_manager_v1::Request::CreateSource { source, .. } => {
                data_init.init(source, CaptureSourceData);
            }
            ext_output_image_capture_source_manager_v1::Request::Destroy => {}
            _ => {}
        }
    }
}

// --- ExtImageCaptureSourceV1 ---

impl<D> Dispatch<ExtImageCaptureSourceV1, CaptureSourceData, D> for ImageCopyCaptureState
where
    D: Dispatch<ExtImageCaptureSourceV1, CaptureSourceData> + 'static,
{
    fn request(
        _state: &mut D,
        _client: &Client,
        _resource: &ExtImageCaptureSourceV1,
        request: ext_image_capture_source_v1::Request,
        _data: &CaptureSourceData,
        _display: &DisplayHandle,
        _data_init: &mut DataInit<'_, D>,
    ) {
        match request {
            ext_image_capture_source_v1::Request::Destroy => {}
            _ => {}
        }
    }
}

// --- ExtImageCopyCaptureManagerV1 (global) ---

impl<D> GlobalDispatch<ExtImageCopyCaptureManagerV1, (), D> for ImageCopyCaptureState
where
    D: GlobalDispatch<ExtImageCopyCaptureManagerV1, ()>
        + Dispatch<ExtImageCopyCaptureManagerV1, ()>
        + Dispatch<ExtImageCopyCaptureSessionV1, CaptureSessionData>
        + Dispatch<ExtImageCopyCaptureCursorSessionV1, CaptureCursorSessionData>
        + Dispatch<ExtImageCopyCaptureFrameV1, CaptureFrameData>
        + Dispatch<ExtImageCaptureSourceV1, CaptureSourceData>
        + ImageCopyCaptureHandler
        + 'static,
{
    fn bind(
        _state: &mut D,
        _display: &DisplayHandle,
        _client: &Client,
        resource: New<ExtImageCopyCaptureManagerV1>,
        _data: &(),
        data_init: &mut DataInit<'_, D>,
    ) {
        data_init.init(resource, ());
        debug!("ext-image-copy-capture-manager bound");
    }
}

impl<D> Dispatch<ExtImageCopyCaptureManagerV1, (), D> for ImageCopyCaptureState
where
    D: Dispatch<ExtImageCopyCaptureManagerV1, ()>
        + Dispatch<ExtImageCopyCaptureSessionV1, CaptureSessionData>
        + Dispatch<ExtImageCopyCaptureCursorSessionV1, CaptureCursorSessionData>
        + Dispatch<ExtImageCopyCaptureFrameV1, CaptureFrameData>
        + Dispatch<ExtImageCaptureSourceV1, CaptureSourceData>
        + ImageCopyCaptureHandler
        + 'static,
{
    fn request(
        _state: &mut D,
        _client: &Client,
        _resource: &ExtImageCopyCaptureManagerV1,
        request: ext_image_copy_capture_manager_v1::Request,
        _data: &(),
        _display: &DisplayHandle,
        data_init: &mut DataInit<'_, D>,
    ) {
        match request {
            ext_image_copy_capture_manager_v1::Request::CreateSession {
                session,
                source: _,
                options: _,
            } => {
                data_init.init(session, CaptureSessionData);
            }
            ext_image_copy_capture_manager_v1::Request::CreatePointerCursorSession {
                session,
                ..
            } => {
                data_init.init(session, CaptureCursorSessionData);
            }
            ext_image_copy_capture_manager_v1::Request::Destroy => {}
            _ => {}
        }
    }
}

// --- ExtImageCopyCaptureSessionV1 ---

impl<D> Dispatch<ExtImageCopyCaptureSessionV1, CaptureSessionData, D> for ImageCopyCaptureState
where
    D: Dispatch<ExtImageCopyCaptureSessionV1, CaptureSessionData>
        + Dispatch<ExtImageCopyCaptureFrameV1, CaptureFrameData>
        + 'static,
{
    fn request(
        _state: &mut D,
        _client: &Client,
        _resource: &ExtImageCopyCaptureSessionV1,
        request: ext_image_copy_capture_session_v1::Request,
        _data: &CaptureSessionData,
        _display: &DisplayHandle,
        data_init: &mut DataInit<'_, D>,
    ) {
        match request {
            ext_image_copy_capture_session_v1::Request::CreateFrame { frame } => {
                data_init.init(frame, CaptureFrameData { buffer: None });
            }
            ext_image_copy_capture_session_v1::Request::Destroy => {}
            _ => {}
        }
    }
}

// --- ExtImageCopyCaptureCursorSessionV1 ---

impl<D> Dispatch<ExtImageCopyCaptureCursorSessionV1, CaptureCursorSessionData, D>
    for ImageCopyCaptureState
where
    D: Dispatch<ExtImageCopyCaptureCursorSessionV1, CaptureCursorSessionData>
        + Dispatch<ExtImageCopyCaptureSessionV1, CaptureSessionData>
        + 'static,
{
    fn request(
        _state: &mut D,
        _client: &Client,
        _resource: &ExtImageCopyCaptureCursorSessionV1,
        request: ext_image_copy_capture_cursor_session_v1::Request,
        _data: &CaptureCursorSessionData,
        _display: &DisplayHandle,
        data_init: &mut DataInit<'_, D>,
    ) {
        match request {
            ext_image_copy_capture_cursor_session_v1::Request::GetCaptureSession { session } => {
                data_init.init(session, CaptureSessionData);
            }
            ext_image_copy_capture_cursor_session_v1::Request::Destroy => {}
            _ => {}
        }
    }
}

// --- ExtImageCopyCaptureFrameV1 ---

impl<D> Dispatch<ExtImageCopyCaptureFrameV1, CaptureFrameData, D> for ImageCopyCaptureState
where
    D: Dispatch<ExtImageCopyCaptureFrameV1, CaptureFrameData>
        + ImageCopyCaptureHandler
        + 'static,
{
    fn request(
        state: &mut D,
        _client: &Client,
        resource: &ExtImageCopyCaptureFrameV1,
        request: ext_image_copy_capture_frame_v1::Request,
        data: &CaptureFrameData,
        _display: &DisplayHandle,
        _data_init: &mut DataInit<'_, D>,
    ) {
        match request {
            ext_image_copy_capture_frame_v1::Request::Destroy => {}
            ext_image_copy_capture_frame_v1::Request::AttachBuffer { buffer } => {
                // Buffer attachment — store for use on Capture.
                // Safety: we need interior mutability but the protocol guarantees
                // AttachBuffer is called before Capture. We use unsafe to set the buffer.
                let data_ptr = data as *const CaptureFrameData as *mut CaptureFrameData;
                unsafe { (*data_ptr).buffer = Some(buffer); }
            }
            ext_image_copy_capture_frame_v1::Request::DamageBuffer { .. } => {
                // Damage region noted — we capture the full frame
            }
            ext_image_copy_capture_frame_v1::Request::Capture => {
                if let Some(buffer) = &data.buffer {
                    state.image_copy_frame(ImageCopyCapture {
                        frame: resource.clone(),
                        buffer: buffer.clone(),
                        submitted: false,
                    });
                } else {
                    resource.failed(ext_image_copy_capture_frame_v1::FailureReason::Unknown);
                }
            }
            _ => {}
        }
    }
}

/// Delegate macro for ext-image-copy-capture protocol.
#[macro_export]
macro_rules! delegate_image_copy_capture {
    ($(@<$( $lt:tt $( : $clt:tt $(+ $dlt:tt )* )? ),+>)? $ty: ty) => {
        smithay::reexports::wayland_server::delegate_global_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            smithay::reexports::wayland_protocols::ext::image_copy_capture::v1::server::ext_image_copy_capture_manager_v1::ExtImageCopyCaptureManagerV1: ()
        ] => $crate::protocols::image_copy_capture::ImageCopyCaptureState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            smithay::reexports::wayland_protocols::ext::image_copy_capture::v1::server::ext_image_copy_capture_manager_v1::ExtImageCopyCaptureManagerV1: ()
        ] => $crate::protocols::image_copy_capture::ImageCopyCaptureState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            smithay::reexports::wayland_protocols::ext::image_copy_capture::v1::server::ext_image_copy_capture_session_v1::ExtImageCopyCaptureSessionV1: $crate::protocols::image_copy_capture::CaptureSessionData
        ] => $crate::protocols::image_copy_capture::ImageCopyCaptureState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            smithay::reexports::wayland_protocols::ext::image_copy_capture::v1::server::ext_image_copy_capture_cursor_session_v1::ExtImageCopyCaptureCursorSessionV1: $crate::protocols::image_copy_capture::CaptureCursorSessionData
        ] => $crate::protocols::image_copy_capture::ImageCopyCaptureState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            smithay::reexports::wayland_protocols::ext::image_copy_capture::v1::server::ext_image_copy_capture_frame_v1::ExtImageCopyCaptureFrameV1: $crate::protocols::image_copy_capture::CaptureFrameData
        ] => $crate::protocols::image_copy_capture::ImageCopyCaptureState);

        smithay::reexports::wayland_server::delegate_global_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            smithay::reexports::wayland_protocols::ext::image_capture_source::v1::server::ext_output_image_capture_source_manager_v1::ExtOutputImageCaptureSourceManagerV1: ()
        ] => $crate::protocols::image_copy_capture::ImageCopyCaptureState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            smithay::reexports::wayland_protocols::ext::image_capture_source::v1::server::ext_output_image_capture_source_manager_v1::ExtOutputImageCaptureSourceManagerV1: ()
        ] => $crate::protocols::image_copy_capture::ImageCopyCaptureState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            smithay::reexports::wayland_protocols::ext::image_capture_source::v1::server::ext_image_capture_source_v1::ExtImageCaptureSourceV1: $crate::protocols::image_copy_capture::CaptureSourceData
        ] => $crate::protocols::image_copy_capture::ImageCopyCaptureState);
    };
}
