//! The screencast stream: negotiation state machine, per-frame feeding and
//! the dmabuf frame handles.

use std::cell::RefCell;
use std::rc::Rc;

use pipewire::core::Core;
use pipewire::spa;
use pipewire::spa::pod::Pod;
use pipewire::stream::{Stream, StreamFlags, StreamListener, StreamState};
use smithay::backend::allocator::dmabuf::Dmabuf;
use smithay::backend::allocator::gbm::GbmDevice;
use smithay::backend::drm::DrmDeviceFd;
use tracing::{debug, info, warn};

use super::buffer::{self, DmaSlot};
use super::format::{build_buffers_param, build_format_params};

/// Negotiated stream state shared between the listener callbacks and the
/// per-frame feed (both on the render thread).
#[derive(Default)]
struct CastShared {
    /// Answers the portal's Start once the daemon assigns the node id — it is
    /// not known at connect time, only when the stream reaches `Paused`.
    reply: Option<async_channel::Sender<Option<u32>>>,
    /// Set once PipeWire fixates the format; `None` until then.
    format: Option<(u32, u32)>,
    /// The negotiated buffers carry dmabuf fds (decided per `add_buffer` from
    /// the actual data-type mask, not from the fixated format): blit into
    /// them. Otherwise the buffers are MemPtr: CPU readback via `feed`.
    buffers_dma: bool,
    /// The consumer fixated a BGR* format; the RGBA readback needs an R/B swap.
    swap_rb: bool,
    /// The fixated format carries alpha; force it opaque (scanout alpha < 255).
    force_alpha: bool,
    /// True between the negotiated `Streaming` state and teardown.
    ready: bool,
    /// The stream errored; the owner tears the cast down on the next frame.
    dead: bool,
    /// Allocates the dmabuf slots from `add_buffer`; `None` = SHM only.
    gbm: Option<GbmDevice<DrmDeviceFd>>,
    dma: Vec<DmaSlot>,
}

/// One active portal screencast: a video-source stream feeding a consumer
/// (OBS, xdg-desktop-portal). Fed the composed frame each vblank.
pub struct Cast {
    stream: Stream,
    _listener: StreamListener<()>,
    shared: Rc<RefCell<CastShared>>,
    /// `None` = the primary; `Some(name)` = an extend/mirror head.
    pub output: Option<String>,
    /// The portal session that owns this cast; only its Close stops it.
    pub session: String,
    pub width: u32,
    pub height: u32,
}

impl Cast {
    pub(super) fn new(
        core: &Core,
        output: Option<String>,
        width: u32,
        height: u32,
        gbm: Option<GbmDevice<DrmDeviceFd>>,
        session: String,
        reply: async_channel::Sender<Option<u32>>,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        use pipewire::properties::properties;

        let stream = Stream::new(
            core,
            "forge-screencast",
            properties! {
                *pipewire::keys::MEDIA_CLASS => "Video/Source",
                *pipewire::keys::MEDIA_CATEGORY => "Capture",
                *pipewire::keys::MEDIA_ROLE => "Screen",
                *pipewire::keys::NODE_DESCRIPTION => "forge screencast",
            },
        )?;

        let dma_capable = gbm.is_some();
        let shared = Rc::new(RefCell::new(CastShared {
            reply: Some(reply),
            gbm,
            ..CastShared::default()
        }));

        let sh_state = shared.clone();
        let sh_param = shared.clone();
        let sh_add = shared.clone();
        let sh_remove = shared.clone();
        let listener = stream
            .add_local_listener::<()>()
            .state_changed(move |stream, _, old, new| {
                debug!("screencast state {old:?} -> {new:?}");
                let mut sh = sh_state.borrow_mut();
                sh.ready = matches!(new, StreamState::Streaming);
                match new {
                    StreamState::Paused | StreamState::Streaming => {
                        if let Some(tx) = sh.reply.take() {
                            let id = stream.node_id();
                            info!("screencast node {id}");
                            let _ = tx.try_send(Some(id));
                        }
                    }
                    StreamState::Error(_) => {
                        sh.dead = true;
                        if let Some(tx) = sh.reply.take() {
                            let _ = tx.try_send(None);
                        }
                    }
                    _ => {}
                }
            })
            .param_changed(move |stream, _, id, param| {
                let Some(param) = param else { return };
                if id != spa::param::ParamType::Format.as_raw() {
                    return;
                }
                let Ok((mt, ms)) = spa::param::format_utils::parse_format(param) else { return };
                if mt != spa::param::format::MediaType::Video
                    || ms != spa::param::format::MediaSubtype::Raw
                {
                    return;
                }
                let mut info = spa::param::video::VideoInfoRaw::default();
                if info.parse(param).is_err() {
                    return;
                }
                let size = info.size();
                let (dt_default, dt_mask) = {
                    use spa::param::video::VideoFormat;
                    let mut sh = sh_param.borrow_mut();
                    sh.format = Some((size.width, size.height));
                    // Buffers are (re)allocated after this; `add_buffer` sets
                    // it again if the negotiated type mask carries DmaBuf.
                    sh.buffers_dma = false;
                    sh.swap_rb = matches!(info.format(), VideoFormat::BGRx | VideoFormat::BGRA);
                    sh.force_alpha = matches!(info.format(), VideoFormat::RGBA | VideoFormat::BGRA);
                    if sh.gbm.is_some() {
                        (
                            1i32 << spa::sys::SPA_DATA_DmaBuf,
                            (1i32 << spa::sys::SPA_DATA_DmaBuf)
                                | (1i32 << spa::sys::SPA_DATA_MemPtr),
                        )
                    } else {
                        (1i32 << spa::sys::SPA_DATA_MemPtr, 1i32 << spa::sys::SPA_DATA_MemPtr)
                    }
                };

                // Announce the buffer layout for the fixated size, offering
                // both carriers when dmabuf is possible — the negotiation
                // picks what the consumer can take. The size comes back from
                // the PipeWire negotiation, so the arithmetic is checked.
                let stride_buf = i32::try_from(size.width)
                    .ok()
                    .and_then(|w| w.checked_mul(4))
                    .and_then(|stride| {
                        i32::try_from(size.height)
                            .ok()
                            .and_then(|h| stride.checked_mul(h))
                            .map(|total| (stride, total))
                    });
                let Some((stride, buf_size)) = stride_buf else {
                    warn!("screencast: absurd negotiated size {}x{}", size.width, size.height);
                    return;
                };
                let buffers = build_buffers_param(buf_size, stride, dt_default, dt_mask);
                let mut params = [Pod::from_bytes(&buffers).unwrap()];
                let _ = stream.update_params(&mut params);
            })
            .add_buffer(move |_stream, _, raw| {
                let mut sh = sh_add.borrow_mut();
                let Some(gbm) = sh.gbm.clone() else { return };
                let Some((w, h)) = sh.format else { return };
                match unsafe { buffer::fill_dma_data(raw, &gbm, w, h) } {
                    Ok(Some(slot)) => {
                        sh.dma.push(slot);
                        sh.buffers_dma = true;
                    }
                    // MemPtr negotiation: PipeWire's own allocation stands.
                    Ok(None) => {}
                    Err(e) => {
                        // The set was negotiated as dmabuf; a partially filled
                        // set cannot be fed coherently. Kill the cast (torn
                        // down on the next frame) rather than serving a mix.
                        warn!("screencast dmabuf alloc failed: {e}");
                        sh.dead = true;
                    }
                }
            })
            .remove_buffer(move |_stream, _, raw| {
                let mut sh = sh_remove.borrow_mut();
                if let Some(fd) = unsafe { buffer::buffer_fd(raw) } {
                    sh.dma.retain(|s| s.fd != fd);
                }
            })
            .register()?;

        let pods = build_format_params(width, height, dma_capable);
        let mut params: Vec<&Pod> = pods.iter().map(|b| Pod::from_bytes(b).unwrap()).collect();
        stream.connect(
            spa::utils::Direction::Output,
            None,
            StreamFlags::DRIVER | StreamFlags::MAP_BUFFERS,
            &mut params,
        )?;

        Ok(Self { stream, _listener: listener, shared, output, session, width, height })
    }

    /// The negotiated stream wants CPU-readback frames.
    pub fn wants_shm(&self) -> bool {
        let sh = self.shared.borrow();
        sh.ready && !sh.buffers_dma
    }

    /// The stream errored and the cast should be torn down.
    pub fn is_dead(&self) -> bool {
        self.shared.borrow().dead
    }

    /// The source changed size (mode or orientation switch): re-announce the
    /// formats so the consumer renegotiates and the buffers are reallocated.
    pub fn renegotiate(&mut self, width: u32, height: u32) {
        if self.width == width && self.height == height {
            return;
        }
        self.width = width;
        self.height = height;
        let dma = self.shared.borrow().gbm.is_some();
        let pods = build_format_params(width, height, dma);
        let mut params: Vec<&Pod> = pods.iter().map(|b| Pod::from_bytes(b).unwrap()).collect();
        let _ = self.stream.update_params(&mut params);
        info!("screencast renegotiating to {width}x{height}");
    }

    /// Dequeue a free dmabuf slot to blit the composed frame into, or `None`
    /// when the stream is not a ready dmabuf cast or no buffer is free.
    pub fn dma_frame(&self) -> Option<DmaFrame<'_>> {
        {
            let sh = self.shared.borrow();
            if !sh.ready || !sh.buffers_dma {
                return None;
            }
        }
        let mut buffer = self.stream.dequeue_buffer()?;
        let fd = buffer.datas_mut().first()?.as_raw().fd;
        let (dmabuf, stride) = {
            let sh = self.shared.borrow();
            let slot = sh.dma.iter().find(|s| s.fd == fd)?;
            (slot.dmabuf.clone(), slot.stride)
        };
        Some(DmaFrame { buffer, dmabuf, stride, height: self.height })
    }

    /// Push one composed frame (tightly-packed RGBA, `stride` bytes/row,
    /// top-down) into a free MemPtr stream buffer, converting to the fixated
    /// format. No-op if the consumer has no buffer free or the format is not
    /// yet negotiated — dropping a screencast frame is fine.
    pub fn feed(&self, pixels: &[u8], stride: usize) {
        let (swap_rb, force_alpha) = {
            let sh = self.shared.borrow();
            if !sh.ready || sh.format.is_none() || sh.buffers_dma {
                return;
            }
            (sh.swap_rb, sh.force_alpha)
        };
        let Some(mut buffer) = self.stream.dequeue_buffer() else { return };
        let datas = buffer.datas_mut();
        let Some(data) = datas.first_mut() else { return };
        let dst_stride = self.width as usize * 4;
        let rows = self.height as usize;

        // Publish the chunk as valid only after the whole frame landed —
        // anything short would hand the consumer stale or uninitialised rows.
        let mut complete = false;
        if let Some(dst) = data.data() {
            if stride >= dst_stride
                && pixels.len() >= stride * rows
                && dst.len() >= dst_stride * rows
            {
                for row in 0..rows {
                    dst[row * dst_stride..][..dst_stride]
                        .copy_from_slice(&pixels[row * stride..][..dst_stride]);
                }
                let filled = dst_stride * rows;
                if swap_rb {
                    for px in dst[..filled].chunks_exact_mut(4) {
                        px.swap(0, 2);
                    }
                }
                if force_alpha {
                    for px in dst[..filled].chunks_exact_mut(4) {
                        px[3] = 255;
                    }
                }
                complete = true;
            }
        }
        let chunk = data.chunk_mut();
        *chunk.offset_mut() = 0;
        *chunk.stride_mut() = dst_stride as i32;
        *chunk.size_mut() = if complete { (dst_stride * rows) as u32 } else { 0 };
    }
}

/// A blit destination handed out by `Cast::dma_frame`. Queued back to the
/// consumer on drop; `finish` marks the chunk valid, `discard` empty.
pub struct DmaFrame<'a> {
    buffer: pipewire::buffer::Buffer<'a>,
    pub dmabuf: Dmabuf,
    stride: i32,
    height: u32,
}

impl DmaFrame<'_> {
    pub fn finish(mut self) {
        let (stride, height) = (self.stride, self.height);
        if let Some(data) = self.buffer.datas_mut().first_mut() {
            let chunk = data.chunk_mut();
            *chunk.offset_mut() = 0;
            *chunk.stride_mut() = stride;
            *chunk.size_mut() = stride as u32 * height;
        }
    }

    pub fn discard(mut self) {
        if let Some(data) = self.buffer.datas_mut().first_mut() {
            *data.chunk_mut().size_mut() = 0;
        }
    }
}
