//! Imitates gamescope's PipeWire producer for local end-to-end tests (spec 9.3).
//!
//! Mirrors `src/pipewire.cpp` of gamescope 3.16.28: node properties, the four EnumFormat
//! objects with the custom `requested_size`/`focus_appid` keys, `calculate_capture_size`,
//! MemFd buffers, CORRUPTED marking of buffers whose size does not match the render size,
//! and renegotiation of the format when the capture size changes.
//!
//! cargo run --release --example fake_gamescope -- --name gamescope --width 3840 --height 2160

use std::cell::RefCell;
use std::io::{Cursor, IsTerminal};
use std::rc::Rc;
use std::time::Duration;

use clap::Parser;
use gamescope_hyperhdr::capture::{
    requested_size_of, SPA_FORMAT_VIDEO_GAMESCOPE_FOCUS_APPID, SPA_FORMAT_VIDEO_REQUESTED_SIZE,
    SPA_META_REQUESTED_SIZE_SCALE,
};
use pipewire as pw;
use pw::properties::properties;
use pw::spa;
use pw::stream::{Stream, StreamFlags, StreamRc, StreamState};
use spa::param::format::{FormatProperties, MediaSubtype, MediaType};
use spa::param::video::{VideoFormat, VideoInfoRaw};
use spa::param::ParamType;
use spa::pod::deserialize::PodDeserializer;
use spa::pod::serialize::PodSerializer;
use spa::pod::{ChoiceValue, Object, Pod, Property, PropertyFlags, Value};
use spa::sys as spa_sys;
use spa::utils::{Choice, ChoiceEnum, ChoiceFlags, Direction, Fraction, Id, Rectangle, SpaTypes};
use tracing::{debug, error, info, warn};

const DRM_FORMAT_MOD_LINEAR: i64 = 0;

#[derive(Parser)]
#[command(about = "Fake gamescope PipeWire producer")]
struct Opt {
    /// node.name of the stream
    #[arg(long, default_value = "gamescope")]
    name: String,
    /// Output (display) width
    #[arg(long, default_value_t = 3840)]
    width: u32,
    /// Output (display) height
    #[arg(long, default_value_t = 2160)]
    height: u32,
    /// Frames per second while streaming
    #[arg(long, default_value_t = 60)]
    fps: u32,
    /// Stop producing frames after the first few correctly sized ones (static picture)
    #[arg(long = "static")]
    static_picture: bool,
    /// Ignore requested_size, like a gamescope without the extension
    #[arg(long)]
    ignore_requested: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Dim {
    w: u32,
    h: u32,
}

#[derive(Debug, Clone, Copy)]
struct Negotiated {
    format: VideoFormat,
    size: Dim,
    stride: u32,
    dmabuf: bool,
}

impl Negotiated {
    /// `stride * height`, plus the chroma plane for NV12.
    fn byte_size(&self, height: u32) -> u32 {
        let mut size = self.stride * height;
        if self.format == VideoFormat::NV12 {
            size += self.stride * height.div_ceil(2);
        }
        size
    }
}

struct Fake {
    output: Dim,
    requested: Dim,
    capture: Dim,
    format: Option<Negotiated>,
    streaming: bool,
    seq: u64,
    frame_no: u64,
    good: u64,
    corrupted: u64,
    out_of_buffers: u64,
    static_picture: bool,
    ignore_requested: bool,
}

impl Fake {
    /// gamescope's `calculate_capture_size`, including its float math.
    fn calculate_capture_size(&mut self) {
        let (ow, oh) = (self.output.w, self.output.h);
        let (rw, rh) = (self.requested.w, self.requested.h);
        let (mut cw, mut ch) = (ow, oh);
        if rw > 0 && rh > 0 && (ow > rw || oh > rh) {
            let ratio_w = rw as f32 / ow as f32;
            let ratio_h = rh as f32 / oh as f32;
            if ratio_w <= ratio_h {
                cw = rw;
                ch = (ratio_w * oh as f32).ceil() as u32;
            } else {
                cw = (ratio_h * ow as f32).ceil() as u32;
                ch = rh;
            }
        }
        self.capture = Dim {
            w: cw.next_multiple_of(2),
            h: ch.next_multiple_of(2),
        };
    }
}

/// A MemFd buffer; `texture` is the size frames are rendered at (gamescope's CVulkanTexture).
struct ShmBuffer {
    fd: i32,
    ptr: *mut u8,
    len: usize,
    video: Negotiated,
    texture: Dim,
}

fn serialize(object: Object) -> Vec<u8> {
    PodSerializer::serialize(Cursor::new(Vec::new()), &Value::Object(object))
        .expect("serializing into a Vec cannot fail")
        .0
        .into_inner()
}

fn id_enum(default: u32, alternatives: &[u32]) -> Value {
    Value::Choice(ChoiceValue::Id(Choice(
        ChoiceFlags::empty(),
        ChoiceEnum::Enum {
            default: Id(default),
            alternatives: alternatives.iter().map(|v| Id(*v)).collect(),
        },
    )))
}

/// gamescope's `build_format_params`: BGRx and NV12, each with and without a modifier.
fn build_format_params(capture: Dim) -> Vec<Vec<u8>> {
    let mut params = Vec::new();
    for format in [VideoFormat::BGRx, VideoFormat::NV12] {
        for with_modifier in [true, false] {
            let mut properties = vec![
                Property::new(
                    FormatProperties::MediaType.as_raw(),
                    Value::Id(Id(MediaType::Video.as_raw())),
                ),
                Property::new(
                    FormatProperties::MediaSubtype.as_raw(),
                    Value::Id(Id(MediaSubtype::Raw.as_raw())),
                ),
                Property::new(
                    FormatProperties::VideoFormat.as_raw(),
                    Value::Id(Id(format.as_raw())),
                ),
                Property::new(
                    FormatProperties::VideoSize.as_raw(),
                    Value::Rectangle(Rectangle {
                        width: capture.w,
                        height: capture.h,
                    }),
                ),
                Property::new(
                    FormatProperties::VideoFramerate.as_raw(),
                    Value::Fraction(Fraction { num: 0, denom: 1 }),
                ),
                Property::new(
                    SPA_FORMAT_VIDEO_REQUESTED_SIZE,
                    Value::Choice(ChoiceValue::Rectangle(Choice(
                        ChoiceFlags::empty(),
                        ChoiceEnum::Range {
                            default: Rectangle {
                                width: 0,
                                height: 0,
                            },
                            min: Rectangle {
                                width: 0,
                                height: 0,
                            },
                            max: Rectangle {
                                width: u32::MAX,
                                height: u32::MAX,
                            },
                        },
                    ))),
                ),
                Property::new(
                    SPA_FORMAT_VIDEO_GAMESCOPE_FOCUS_APPID,
                    Value::Choice(ChoiceValue::Long(Choice(
                        ChoiceFlags::empty(),
                        ChoiceEnum::Range {
                            default: 0,
                            min: i64::MIN,
                            max: i64::MAX,
                        },
                    ))),
                ),
            ];
            if format == VideoFormat::NV12 {
                properties.push(Property::new(
                    FormatProperties::VideoColorMatrix.as_raw(),
                    id_enum(
                        spa_sys::SPA_VIDEO_COLOR_MATRIX_BT601,
                        &[
                            spa_sys::SPA_VIDEO_COLOR_MATRIX_BT601,
                            spa_sys::SPA_VIDEO_COLOR_MATRIX_BT709,
                        ],
                    ),
                ));
                properties.push(Property::new(
                    FormatProperties::VideoColorRange.as_raw(),
                    id_enum(
                        spa_sys::SPA_VIDEO_COLOR_RANGE_16_235,
                        &[
                            spa_sys::SPA_VIDEO_COLOR_RANGE_16_235,
                            spa_sys::SPA_VIDEO_COLOR_RANGE_0_255,
                        ],
                    ),
                ));
            }
            if with_modifier {
                properties.push(Property {
                    key: FormatProperties::VideoModifier.as_raw(),
                    flags: PropertyFlags::MANDATORY,
                    value: Value::Choice(ChoiceValue::Long(Choice(
                        ChoiceFlags::empty(),
                        ChoiceEnum::Enum {
                            default: DRM_FORMAT_MOD_LINEAR,
                            alternatives: vec![DRM_FORMAT_MOD_LINEAR],
                        },
                    ))),
                });
            }
            params.push(serialize(Object {
                type_: SpaTypes::ObjectParamFormat.as_raw(),
                id: ParamType::EnumFormat.as_raw(),
                properties,
            }));
        }
    }
    params
}

fn meta_param(meta_type: u32, size: usize) -> Vec<u8> {
    serialize(Object {
        type_: SpaTypes::ObjectParamMeta.as_raw(),
        id: ParamType::Meta.as_raw(),
        properties: vec![
            Property::new(spa_sys::SPA_PARAM_META_type, Value::Id(Id(meta_type))),
            Property::new(spa_sys::SPA_PARAM_META_size, Value::Int(size as i32)),
        ],
    })
}

/// Buffers and Meta params set by `stream_handle_param_changed`.
fn buffer_params(n: &Negotiated) -> Vec<Vec<u8>> {
    let data_type = if n.dmabuf {
        1 << spa_sys::SPA_DATA_DmaBuf
    } else {
        1 << spa_sys::SPA_DATA_MemFd
    };
    let buffers = serialize(Object {
        type_: SpaTypes::ObjectParamBuffers.as_raw(),
        id: ParamType::Buffers.as_raw(),
        properties: vec![
            Property::new(
                spa_sys::SPA_PARAM_BUFFERS_buffers,
                Value::Choice(ChoiceValue::Int(Choice(
                    ChoiceFlags::empty(),
                    ChoiceEnum::Range {
                        default: 4,
                        min: 1,
                        max: 8,
                    },
                ))),
            ),
            Property::new(spa_sys::SPA_PARAM_BUFFERS_blocks, Value::Int(1)),
            Property::new(
                spa_sys::SPA_PARAM_BUFFERS_size,
                Value::Int(n.byte_size(n.size.h) as i32),
            ),
            Property::new(
                spa_sys::SPA_PARAM_BUFFERS_stride,
                Value::Int(n.stride as i32),
            ),
            Property::new(
                spa_sys::SPA_PARAM_BUFFERS_dataType,
                Value::Choice(ChoiceValue::Int(Choice(
                    ChoiceFlags::empty(),
                    ChoiceEnum::Flags {
                        default: data_type,
                        flags: vec![],
                    },
                ))),
            ),
        ],
    });
    vec![
        buffers,
        meta_param(
            spa_sys::SPA_META_Header,
            size_of::<spa_sys::spa_meta_header>(),
        ),
        meta_param(SPA_META_REQUESTED_SIZE_SCALE, size_of::<f32>()),
    ]
}

fn update_params(stream: &Stream, pods: &[Vec<u8>]) {
    let mut params: Vec<&Pod> = pods
        .iter()
        .map(|p| Pod::from_bytes(p).expect("valid pod"))
        .collect();
    if let Err(e) = stream.update_params(&mut params) {
        error!("pw_stream_update_params failed: {e}");
    }
}

fn has_modifier(format: &Pod) -> bool {
    matches!(
        PodDeserializer::deserialize_any_from(format.as_bytes()),
        Ok((_, Value::Object(o))) if o.properties.iter().any(|p| p.key == FormatProperties::VideoModifier.as_raw())
    )
}

fn on_param_changed(stream: &Stream, fake: &RefCell<Fake>, id: u32, param: Option<&Pod>) {
    if id != ParamType::Format.as_raw() {
        return;
    }
    let Some(param) = param else {
        return;
    };
    let mut info = VideoInfoRaw::new();
    if info.parse(param).is_err() {
        error!("spa_format_video_raw_parse failed");
        return;
    }
    let requested = requested_size_of(param).map_or(Dim { w: 0, h: 0 }, |r| Dim {
        w: r.width,
        h: r.height,
    });
    let negotiated = {
        let mut f = fake.borrow_mut();
        f.requested = if f.ignore_requested {
            Dim { w: 0, h: 0 }
        } else {
            requested
        };
        f.calculate_capture_size();
        let size = Dim {
            w: info.size().width,
            h: info.size().height,
        };
        let bpp = if info.format() == VideoFormat::NV12 {
            1
        } else {
            4
        };
        let negotiated = Negotiated {
            format: info.format(),
            size,
            stride: (size.w * bpp).next_multiple_of(4),
            dmabuf: has_modifier(param),
        };
        f.format = Some(negotiated);
        info!(
            "format changed (size {}x{}, requested {}x{}, capture {}x{}, format {:?}, stride {}, dmabuf {})",
            size.w,
            size.h,
            requested.w,
            requested.h,
            f.capture.w,
            f.capture.h,
            negotiated.format,
            negotiated.stride,
            negotiated.dmabuf
        );
        negotiated
    };
    update_params(stream, &buffer_params(&negotiated));
}

/// # Safety
/// `pw_buffer` must be the buffer passed to the `add_buffer` event.
unsafe fn on_add_buffer(fake: &Fake, pw_buffer: *mut pw::sys::pw_buffer) {
    let spa_buffer = (*pw_buffer).buffer;
    if spa_buffer.is_null() || (*spa_buffer).n_datas < 1 {
        return;
    }
    let data = &mut *(*spa_buffer).datas;
    let Some(n) = fake.format else {
        error!("add_buffer without a format");
        return;
    };
    if data.type_ & (1 << spa_sys::SPA_DATA_MemFd) == 0 {
        error!(
            "unsupported data types {:#x}: the fake only implements MemFd",
            data.type_
        );
        data.type_ = spa_sys::SPA_DATA_Invalid;
        return;
    }
    let len = n.byte_size(n.size.h) as usize;
    let fd = libc::memfd_create(c"gamescope-fake-pw".as_ptr(), libc::MFD_CLOEXEC);
    if fd < 0 {
        error!("memfd_create failed");
        return;
    }
    if libc::ftruncate(fd, len as libc::off_t) != 0 {
        error!("ftruncate failed");
        libc::close(fd);
        return;
    }
    let ptr = libc::mmap(
        std::ptr::null_mut(),
        len,
        libc::PROT_READ | libc::PROT_WRITE,
        libc::MAP_SHARED,
        fd,
        0,
    );
    if ptr == libc::MAP_FAILED {
        error!("mmap failed");
        libc::close(fd);
        return;
    }
    data.type_ = spa_sys::SPA_DATA_MemFd;
    data.flags = spa_sys::SPA_DATA_FLAG_READABLE;
    data.fd = i64::from(fd);
    data.mapoffset = 0;
    data.maxsize = len as u32;
    data.data = ptr;
    debug!(
        len,
        "added buffer, texture {}x{}", fake.capture.w, fake.capture.h
    );

    let buffer = Box::new(ShmBuffer {
        fd,
        ptr: ptr.cast(),
        len,
        video: n,
        texture: fake.capture,
    });
    (*pw_buffer).user_data = Box::into_raw(buffer).cast();
}

/// # Safety
/// `pw_buffer` must be the buffer passed to the `remove_buffer` event.
unsafe fn on_remove_buffer(pw_buffer: *mut pw::sys::pw_buffer) {
    let ptr = (*pw_buffer).user_data.cast::<ShmBuffer>();
    if ptr.is_null() {
        return;
    }
    (*pw_buffer).user_data = std::ptr::null_mut();
    let buffer = Box::from_raw(ptr);
    libc::munmap(buffer.ptr.cast(), buffer.len);
    libc::close(buffer.fd);
}

/// RGB colour bars that move by one bar every 8 frames.
fn render(buffer: &mut ShmBuffer, frame_no: u64) {
    const BARS: [[u8; 3]; 8] = [
        [235, 235, 235],
        [235, 235, 16],
        [16, 235, 235],
        [16, 235, 16],
        [235, 16, 235],
        [235, 16, 16],
        [16, 16, 235],
        [16, 16, 16],
    ];
    // SAFETY: `ptr` maps `len` bytes for the lifetime of the buffer.
    let mem = unsafe { std::slice::from_raw_parts_mut(buffer.ptr, buffer.len) };
    let (w, h) = (buffer.texture.w as usize, buffer.texture.h as usize);
    let stride = buffer.video.stride as usize;
    if buffer.video.format == VideoFormat::NV12 {
        mem.fill(128);
        return;
    }
    let shift = (frame_no / 8) as usize;
    for x in 0..w {
        let [r, g, b] = BARS[(x * BARS.len() / w + shift) % BARS.len()];
        mem[x * 4..x * 4 + 4].copy_from_slice(&[b, g, r, 0xFF]);
    }
    for y in 1..h {
        mem.copy_within(0..w * 4, y * stride);
    }
}

/// gamescope's `paint_pipewire` + `dispatch_nudge` + `copy_buffer`, once per frame.
fn tick(stream: &Stream, fake: &RefCell<Fake>) {
    let renegotiate = {
        let f = fake.borrow();
        if !f.streaming {
            return;
        }
        let current = f.format.map_or(Dim { w: 0, h: 0 }, |n| n.size);
        (f.capture != current).then(|| {
            debug!(
                "renegotiating stream params (size: {}x{})",
                f.capture.w, f.capture.h
            );
            build_format_params(f.capture)
        })
    };
    if let Some(pods) = renegotiate {
        update_params(stream, &pods);
    }

    let mut f = fake.borrow_mut();
    if f.static_picture && f.good >= 5 {
        return;
    }
    // SAFETY: the buffer is queued back below and only accessed in between.
    unsafe {
        let raw = stream.dequeue_raw_buffer();
        if raw.is_null() {
            f.out_of_buffers += 1;
            if f.out_of_buffers.is_power_of_two() {
                warn!("out of buffers ({} times)", f.out_of_buffers);
            }
            return;
        }
        let shm = (*raw).user_data.cast::<ShmBuffer>();
        if shm.is_null() {
            stream.queue_raw_buffer(raw);
            return;
        }
        let shm = &mut *shm;
        let spa_buffer = (*raw).buffer;
        let needs_reneg = shm.video.size != shm.texture;

        let header = spa_sys::spa_buffer_find_meta_data(
            spa_buffer,
            spa_sys::SPA_META_Header,
            size_of::<spa_sys::spa_meta_header>(),
        )
        .cast::<spa_sys::spa_meta_header>();
        if !header.is_null() {
            (*header).pts = -1;
            (*header).flags = if needs_reneg {
                spa_sys::SPA_META_HEADER_FLAG_CORRUPTED
            } else {
                0
            };
            (*header).seq = f.seq;
            (*header).dts_offset = 0;
        }
        f.seq += 1;
        let scale = spa_sys::spa_buffer_find_meta_data(
            spa_buffer,
            SPA_META_REQUESTED_SIZE_SCALE,
            size_of::<f32>(),
        )
        .cast::<f32>();
        if !scale.is_null() {
            *scale = shm.texture.w as f32 / f.output.w as f32;
        }

        let data = &mut *(*spa_buffer).datas;
        let chunk = &mut *data.chunk;
        chunk.flags = if needs_reneg {
            spa_sys::SPA_CHUNK_FLAG_CORRUPTED as i32
        } else {
            0
        };
        // Like gamescope, the chunk size follows the current format, the stride the buffer.
        let current = f.format.unwrap_or(shm.video);
        chunk.offset = 0;
        chunk.size = Negotiated {
            stride: shm.video.stride,
            ..current
        }
        .byte_size(current.size.h);
        chunk.stride = shm.video.stride as i32;

        if needs_reneg {
            f.corrupted += 1;
        } else {
            render(shm, f.frame_no);
            f.good += 1;
        }
        f.frame_no += 1;
        drop(f);
        stream.queue_raw_buffer(raw);
    }
}

fn main() -> anyhow::Result<()> {
    let opt = Opt::parse();
    tracing_subscriber::fmt()
        .with_target(false)
        .with_writer(std::io::stderr)
        .with_ansi(std::io::stderr().is_terminal())
        .init();

    pw::init();
    let mainloop = pw::main_loop::MainLoopRc::new(None)?;
    let context = pw::context::ContextRc::new(&mainloop, None)?;
    let core = context.connect_rc(None)?;

    let _signals = [pw::loop_::Signal::INT, pw::loop_::Signal::TERM].map(|signal| {
        let mainloop_weak = mainloop.downgrade();
        mainloop.loop_().add_signal_local(signal, move || {
            if let Some(mainloop) = mainloop_weak.upgrade() {
                mainloop.quit();
            }
        })
    });
    let _core_listener = {
        let mainloop_weak = mainloop.downgrade();
        core.add_listener_local()
            .error(move |id, _, res, message| {
                error!(id, res, "PipeWire error: {message}");
                if id == pw::sys::PW_ID_CORE {
                    if let Some(mainloop) = mainloop_weak.upgrade() {
                        mainloop.quit();
                    }
                }
            })
            .register()
    };

    let fake = Rc::new(RefCell::new(Fake {
        output: Dim {
            w: opt.width,
            h: opt.height,
        },
        requested: Dim { w: 0, h: 0 },
        capture: Dim { w: 0, h: 0 },
        format: None,
        streaming: false,
        seq: 0,
        frame_no: 0,
        good: 0,
        corrupted: 0,
        out_of_buffers: 0,
        static_picture: opt.static_picture,
        ignore_requested: opt.ignore_requested,
    }));
    fake.borrow_mut().calculate_capture_size();

    // gamescope passes only media.class. With PipeWire 1.6 the stream name becomes media.name
    // and node.name defaults to the program name, which is "gamescope" for the real binary,
    // so the fake sets node.name explicitly.
    let stream = StreamRc::new(
        core.clone(),
        &opt.name,
        properties! {
            *pw::keys::MEDIA_CLASS => "Video/Source",
            *pw::keys::NODE_NAME => opt.name.as_str(),
        },
    )?;
    let _listener = {
        let mainloop_weak = mainloop.downgrade();
        stream
            .add_local_listener_with_user_data(fake.clone())
            .state_changed(move |stream, fake, old, new| {
                info!("stream state changed: {old:?} -> {new:?}");
                let mut f = fake.borrow_mut();
                match new {
                    StreamState::Paused => {
                        info!("stream available on node ID: {}", stream.node_id());
                        f.streaming = false;
                        f.seq = 0;
                    }
                    StreamState::Streaming => f.streaming = true,
                    StreamState::Error(_) | StreamState::Unconnected => {
                        if let Some(mainloop) = mainloop_weak.upgrade() {
                            mainloop.quit();
                        }
                    }
                    StreamState::Connecting => {}
                }
            })
            .param_changed(|stream, fake, id, param| on_param_changed(stream, fake, id, param))
            // SAFETY: the pointers come straight from the stream events.
            .add_buffer(|_, fake, buffer| unsafe { on_add_buffer(&fake.borrow(), buffer) })
            .remove_buffer(|_, _, buffer| unsafe { on_remove_buffer(buffer) })
            .register()?
    };

    let pods = build_format_params(fake.borrow().capture);
    let mut params: Vec<&Pod> = pods
        .iter()
        .map(|p| Pod::from_bytes(p).expect("valid pod"))
        .collect();
    stream.connect(
        Direction::Output,
        None,
        StreamFlags::DRIVER | StreamFlags::ALLOC_BUFFERS,
        &mut params,
    )?;

    let frame_timer = {
        let (stream, fake) = (stream.clone(), fake.clone());
        mainloop.loop_().add_timer(move |_| tick(&stream, &fake))
    };
    let interval = Duration::from_secs(1) / opt.fps.max(1);
    let _ = frame_timer.update_timer(Some(interval), Some(interval));

    let report_timer = {
        let fake = fake.clone();
        mainloop.loop_().add_timer(move |_| {
            let f = fake.borrow();
            info!(
                "frames: {} ok, {} corrupted, out of buffers {}, streaming {}, capture {}x{}",
                f.good, f.corrupted, f.out_of_buffers, f.streaming, f.capture.w, f.capture.h
            );
        })
    };
    let _ = report_timer.update_timer(Some(Duration::from_secs(5)), Some(Duration::from_secs(5)));

    info!(
        "fake gamescope '{}' output {}x{} at {} fps{}{}",
        opt.name,
        opt.width,
        opt.height,
        opt.fps,
        if opt.static_picture { ", static" } else { "" },
        if opt.ignore_requested {
            ", ignoring requested_size"
        } else {
            ""
        }
    );
    mainloop.run();
    Ok(())
}