//! Bounded, cancellable Wayland capture. No screenshots or external capture tools.
use memmap2::MmapMut;
use serde::{Deserialize, Serialize};
use std::{
    fs::File,
    os::fd::{AsFd, AsRawFd, FromRawFd},
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    time::{Duration, Instant},
};
use wayland_client::{
    Connection, Dispatch, EventQueue, QueueHandle, WEnum,
    protocol::{wl_buffer, wl_output, wl_registry, wl_shm, wl_shm_pool},
};
use wayland_protocols_wlr::screencopy::v1::client::{
    zwlr_screencopy_frame_v1 as frame, zwlr_screencopy_manager_v1 as manager,
};

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Source {
    pub name: String,
    pub label: String,
    pub width: i32,
    pub height: i32,
    pub x: i32,
    pub y: i32,
    pub scale: i32,
}

struct Output {
    global: u32,
    proxy: wl_output::WlOutput,
    source: Source,
    transform: wl_output::Transform,
}

pub struct Frame {
    pub width: u32,
    pub height: u32,
    pub stride: u32,
    pub pixels: Vec<u8>,
}

struct Buffer {
    _file: File,
    map: MmapMut,
    proxy: wl_buffer::WlBuffer,
    width: u32,
    height: u32,
    stride: u32,
    format: wl_shm::Format,
}

impl Drop for Buffer {
    fn drop(&mut self) {
        self.proxy.destroy();
    }
}

#[derive(Default)]
struct State {
    outputs: Vec<Output>,
    manager: Option<manager::ZwlrScreencopyManagerV1>,
    shm: Option<wl_shm::WlShm>,
    buffer: Option<Buffer>,
    result: Option<Result<(), String>>,
    flipped: bool,
}

pub struct Capture {
    connection: Connection,
    queue: EventQueue<State>,
    state: State,
}

impl Capture {
    pub fn connect(stop: &AtomicBool) -> Result<Self, String> {
        let connection = Connection::connect_to_env().map_err(|e| e.to_string())?;
        let queue = connection.new_event_queue();
        connection.display().get_registry(&queue.handle(), ());
        let mut this = Self {
            connection,
            queue,
            state: State::default(),
        };
        // Two syncs deliver registry bindings and output metadata without an
        // uninterruptible roundtrip on a stalled compositor.
        this.sync(stop)?;
        this.sync(stop)?;
        if this.state.manager.is_none() || this.state.shm.is_none() {
            return Err("Ferese screen capture is unavailable in this session".into());
        }
        Ok(this)
    }

    fn sync(&mut self, stop: &AtomicBool) -> Result<(), String> {
        let done = Arc::new(AtomicBool::new(false));
        self.connection
            .display()
            .sync(&self.queue.handle(), done.clone());
        let deadline = Instant::now() + Duration::from_secs(5);
        while !done.load(Ordering::Relaxed) {
            self.dispatch(stop, deadline)?;
        }
        Ok(())
    }

    pub fn sources(&self) -> Vec<Source> {
        self.state
            .outputs
            .iter()
            .map(|o| {
                let mut source = o.source.clone();
                if swaps_axes(o.transform) {
                    std::mem::swap(&mut source.width, &mut source.height);
                }
                source
            })
            .collect()
    }

    pub fn frame(
        &mut self,
        name: &str,
        cursor: bool,
        stop: &AtomicBool,
        mut pixels: Vec<u8>,
    ) -> Result<Frame, String> {
        let output = self
            .state
            .outputs
            .iter()
            .find(|o| o.source.name == name)
            .ok_or("Shared monitor disconnected")?;

        let transform = output.transform;
        self.state.result = None;
        self.state.flipped = false;

        let frame = self.state.manager.as_ref().unwrap().capture_output(
            cursor as i32,
            &output.proxy,
            &self.queue.handle(),
            (),
        );
        let deadline = Instant::now() + Duration::from_secs(5);

        while self.state.result.is_none() {
            if let Err(error) = self.dispatch(stop, deadline) {
                frame.destroy();
                return Err(error);
            }
        }

        frame.destroy();
        self.state.result.take().unwrap()?;

        let buffer = self
            .state
            .buffer
            .as_ref()
            .ok_or("Capture supplied no buffer")?;
        let (width, height) = copy_oriented_pixels(
            &buffer.map,
            buffer.width,
            buffer.height,
            buffer.stride,
            transform,
            self.state.flipped,
            &mut pixels,
        );

        Ok(Frame {
            width,
            height,
            stride: width * 4,
            pixels,
        })
    }

    fn dispatch(&mut self, stop: &AtomicBool, deadline: Instant) -> Result<(), String> {
        if stop.load(Ordering::Relaxed) {
            return Err("Capture cancelled".into());
        }

        if Instant::now() >= deadline {
            return Err("Compositor capture timed out".into());
        }

        self.queue
            .dispatch_pending(&mut self.state)
            .map_err(|e| e.to_string())?;
        self.connection.flush().map_err(|e| e.to_string())?;

        if let Some(read) = self.queue.prepare_read() {
            let mut fd = libc::pollfd {
                fd: read.connection_fd().as_raw_fd(),
                events: libc::POLLIN,
                revents: 0,
            };
            // SAFETY: a single valid pollfd, bounded timeout; Wayland owns the FD.
            let ready = unsafe { libc::poll(&mut fd, 1, 20) };

            if ready > 0 {
                read.read().map_err(|e| e.to_string())?;
            } else if ready < 0
                && std::io::Error::last_os_error().kind() != std::io::ErrorKind::Interrupted
            {
                return Err(std::io::Error::last_os_error().to_string());
            }
        }

        self.queue
            .dispatch_pending(&mut self.state)
            .map_err(|e| e.to_string())?;

        Ok(())
    }
}

impl Dispatch<wl_registry::WlRegistry, ()> for State {
    fn event(
        state: &mut Self,
        registry: &wl_registry::WlRegistry,
        event: wl_registry::Event,
        _: &(),
        _: &Connection,
        qh: &QueueHandle<Self>,
    ) {
        match event {
            wl_registry::Event::Global {
                name,
                interface,
                version,
            } => match interface.as_str() {
                "wl_output" => state.outputs.push(Output {
                    global: name,
                    transform: wl_output::Transform::Normal,
                    proxy: registry.bind(name, version.min(4), qh, name),
                    source: Source {
                        name: format!("output-{name}"),
                        label: "Display".into(),
                        width: 0,
                        height: 0,
                        x: 0,
                        y: 0,
                        scale: 1,
                    },
                }),
                "wl_shm" => state.shm = Some(registry.bind(name, 1, qh, ())),
                "zwlr_screencopy_manager_v1" => {
                    state.manager = Some(registry.bind(name, version.min(3), qh, ()))
                }
                _ => (),
            },
            wl_registry::Event::GlobalRemove { name } => state.outputs.retain(|o| o.global != name),
            _ => (),
        }
    }
}

impl Dispatch<wl_output::WlOutput, u32> for State {
    fn event(
        state: &mut Self,
        _: &wl_output::WlOutput,
        event: wl_output::Event,
        name: &u32,
        _: &Connection,
        _: &QueueHandle<Self>,
    ) {
        let Some(output) = state.outputs.iter_mut().find(|o| o.global == *name) else {
            return;
        };

        match event {
            wl_output::Event::Name { name } => output.source.name = name,
            wl_output::Event::Description { description } => output.source.label = description,
            wl_output::Event::Geometry {
                x, y, transform, ..
            } => {
                if let WEnum::Value(transform) = transform {
                    output.transform = transform;
                }
                output.source.x = x;
                output.source.y = y;
            }
            wl_output::Event::Scale { factor } => output.source.scale = factor.max(1),
            wl_output::Event::Mode {
                flags: WEnum::Value(flags),
                width,
                height,
                ..
            } if flags.contains(wl_output::Mode::Current) => {
                output.source.width = width;
                output.source.height = height;
            }
            _ => (),
        }
    }
}

impl Dispatch<frame::ZwlrScreencopyFrameV1, ()> for State {
    fn event(
        state: &mut Self,
        frame: &frame::ZwlrScreencopyFrameV1,
        event: frame::Event,
        _: &(),
        _: &Connection,
        qh: &QueueHandle<Self>,
    ) {
        match event {
            frame::Event::Buffer {
                format: WEnum::Value(format),
                width,
                height,
                stride,
            } => {
                let result = (|| {
                    if format != wl_shm::Format::Argb8888 && format != wl_shm::Format::Xrgb8888 {
                        return Err("Unsupported capture pixel format".into());
                    }

                    let bytes = buffer_size(width, height, stride)?;
                    if let Some(buffer) = &state.buffer
                        && (buffer.width, buffer.height, buffer.stride, buffer.format)
                            == (width, height, stride, format)
                    {
                        frame.copy(&buffer.proxy);
                        return Ok(());
                    }
                    // SAFETY: fixed NUL-terminated name and a newly owned memfd.
                    let fd = unsafe {
                        libc::memfd_create(c"ferese-screencast".as_ptr(), libc::MFD_CLOEXEC)
                    };

                    if fd < 0 {
                        return Err(std::io::Error::last_os_error().to_string());
                    }

                    let file = unsafe { File::from_raw_fd(fd) };
                    file.set_len(bytes as u64).map_err(|e| e.to_string())?;
                    // SAFETY: the file has the checked mapping length and lives with the mapping.
                    let map = unsafe { MmapMut::map_mut(&file) }.map_err(|e| e.to_string())?;
                    let pool = state
                        .shm
                        .as_ref()
                        .ok_or("Missing shared memory global")?
                        .create_pool(file.as_fd(), bytes as i32, qh, ());
                    let proxy = pool.create_buffer(
                        0,
                        width as i32,
                        height as i32,
                        stride as i32,
                        format,
                        qh,
                        (),
                    );

                    pool.destroy();
                    frame.copy(&proxy);
                    state.buffer = Some(Buffer {
                        _file: file,
                        map,
                        proxy,
                        width,
                        height,
                        stride,
                        format,
                    });

                    Ok::<(), String>(())
                })();

                if let Err(e) = result {
                    state.result = Some(Err(e));
                }
            }
            frame::Event::Ready { .. } => state.result = Some(Ok(())),
            frame::Event::Failed => {
                state.result = Some(Err(
                    "Capture ended: session locked, monitor removed, or compositor unavailable"
                        .into(),
                ))
            }
            frame::Event::Flags {
                flags: WEnum::Value(flags),
            } => state.flipped = flags.contains(frame::Flags::YInvert),
            _ => (),
        }
    }
}

fn swaps_axes(transform: wl_output::Transform) -> bool {
    matches!(
        transform,
        wl_output::Transform::_90
            | wl_output::Transform::_270
            | wl_output::Transform::Flipped90
            | wl_output::Transform::Flipped270
    )
}

// Screencopy supplies output-buffer coordinates. Normalize both the output's
// presentation transform and the independent YInvert storage flag before
// publishing upright, tightly packed BGRx frames to PipeWire.
fn copy_oriented_pixels(
    source: &[u8],
    width: u32,
    height: u32,
    stride: u32,
    transform: wl_output::Transform,
    y_invert: bool,
    pixels: &mut Vec<u8>,
) -> (u32, u32) {
    use wl_output::Transform;
    let (out_width, out_height) = if swaps_axes(transform) {
        (height, width)
    } else {
        (width, height)
    };
    let row = out_width as usize * 4;
    pixels.resize(row * out_height as usize, 0);

    // The usual DRM and nested outputs need only one contiguous copy per row.
    if matches!(transform, Transform::Normal | Transform::Flipped180) {
        let flip = y_invert ^ (transform == Transform::Flipped180);
        for (y, dst) in pixels.chunks_exact_mut(row).enumerate() {
            let sy = if flip { height as usize - 1 - y } else { y };
            let start = sy * stride as usize;
            dst.copy_from_slice(&source[start..start + row]);
        }
        return (out_width, out_height);
    }

    for y in 0..out_height {
        for x in 0..out_width {
            let (sx, mut sy) = match transform {
                Transform::_90 => (y, height - 1 - x),
                Transform::_180 => (width - 1 - x, height - 1 - y),
                Transform::_270 => (width - 1 - y, x),
                Transform::Flipped => (width - 1 - x, y),
                Transform::Flipped90 => (y, x),
                Transform::Flipped270 => (width - 1 - y, height - 1 - x),
                _ => (x, y),
            };
            if y_invert {
                sy = height - 1 - sy;
            }
            let src = sy as usize * stride as usize + sx as usize * 4;
            let dst = y as usize * row + x as usize * 4;
            pixels[dst..dst + 4].copy_from_slice(&source[src..src + 4]);
        }
    }
    (out_width, out_height)
}

fn buffer_size(width: u32, height: u32, stride: u32) -> Result<usize, String> {
    let row = width.checked_mul(4).ok_or("Capture dimensions overflow")?;
    let bytes = stride
        .checked_mul(height)
        .ok_or("Capture dimensions overflow")?;

    if width == 0 || height == 0 || stride < row || bytes > 128 * 1024 * 1024 {
        return Err("Unsupported capture dimensions".into());
    }

    Ok(bytes as usize)
}

impl Dispatch<wayland_client::protocol::wl_callback::WlCallback, Arc<AtomicBool>> for State {
    fn event(
        _: &mut Self,
        _: &wayland_client::protocol::wl_callback::WlCallback,
        _: wayland_client::protocol::wl_callback::Event,
        done: &Arc<AtomicBool>,
        _: &Connection,
        _: &QueueHandle<Self>,
    ) {
        done.store(true, Ordering::Relaxed);
    }
}

wayland_client::delegate_noop!(State: ignore wl_shm::WlShm);
wayland_client::delegate_noop!(State: ignore wl_shm_pool::WlShmPool);
wayland_client::delegate_noop!(State: ignore wl_buffer::WlBuffer);
wayland_client::delegate_noop!(State: ignore manager::ZwlrScreencopyManagerV1);

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn capture_orientation_matches_presentation_for_all_transforms() {
        use wl_output::Transform::*;
        // Asymmetric 3x2 image with row padding and distinct four-byte pixels.
        let cases = [
            (Normal, (3, 2), "abcdef"),
            (_90, (2, 3), "daebfc"),
            (_180, (3, 2), "fedcba"),
            (_270, (2, 3), "cfbead"),
            (Flipped, (3, 2), "cbafed"),
            (Flipped90, (2, 3), "adbecf"),
            (Flipped180, (3, 2), "defabc"),
            (Flipped270, (2, 3), "fcebda"),
        ];
        let pixel = |value: u8| [value, value + 1, value + 2, 255];
        let mut pixels = vec![0; 100];
        for (transform, size, expected) in cases {
            for y_invert in [false, true] {
                let mut source = Vec::new();
                let rows = if y_invert {
                    [b"def", b"abc"]
                } else {
                    [b"abc", b"def"]
                };
                for row in rows {
                    for &value in row {
                        source.extend_from_slice(&pixel(value));
                    }
                    source.extend_from_slice(&[0; 4]);
                }
                assert_eq!(
                    copy_oriented_pixels(&source, 3, 2, 16, transform, y_invert, &mut pixels),
                    size,
                );
                let expected: Vec<_> = expected.bytes().flat_map(pixel).collect();
                assert_eq!(pixels, expected, "{transform:?}, YInvert={y_invert}");
            }
        }
    }

    #[test]
    fn rejects_unbounded_and_invalid_buffers() {
        assert_eq!(buffer_size(640, 480, 2560).unwrap(), 1228800);
        assert!(buffer_size(640, 480, 12).is_err());
        assert!(buffer_size(0, 480, 0).is_err());
        assert!(buffer_size(u32::MAX, 2, u32::MAX).is_err());
        assert!(buffer_size(20000, 20000, 80000).is_err());
    }
}