path: root/src/bin/tsv_video_buffer.rs

use std::io::Read;
use std::process::{Command, Stdio};

use ash::vk;
use wgsl_view::{ffmpeg, gpu, hap};

fn is_hap_codec(codec_tag: &str) -> Option<hap::HapFormat> {
    match codec_tag {
        "Hap1" => Some(hap::HapFormat::Bc1),
        "Hap5" => Some(hap::HapFormat::Bc3),
        _ => None,
    }
}

fn main() {
    env_logger::init();

    let args: Vec<String> = std::env::args().skip(1).collect();
    let sep = args.iter().position(|a| a == "--");

    let (own_args, ff_args) = match sep {
        Some(i) => (&args[..i], &args[i + 1..]),
        None => {
            eprintln!("usage: tsv-video-buffer [--name NAME] [--array] [--frames N] -- <ffmpeg input args>");
            std::process::exit(1);
        }
    };

    let mut name = "tsv-video-buffer".to_string();
    let mut use_array = false;
    let mut max_frames: Option<u32> = None;

    let mut i = 0;
    while i < own_args.len() {
        match own_args[i].as_str() {
            "--name" => {
                name = own_args[i + 1].clone();
                i += 2;
            }
            "--array" => {
                use_array = true;
                i += 1;
            }
            "--frames" => {
                max_frames = Some(own_args[i + 1].parse().expect("parse --frames"));
                i += 2;
            }
            other => panic!("unknown argument: {other}"),
        }
    }

    let mut info = ffmpeg::probe_video(ff_args, max_frames);
    let num_frames = match (info.num_frames, max_frames) {
        (Some(a), Some(b)) => u32::min(a, b),
        (Some(a), None) => a,
        (None, Some(b)) => b,
        _ => panic!("failed to detect duration, please specify --frames"),
    };
    info.num_frames = Some(num_frames);
    let hap_format = is_hap_codec(&info.codec_tag);

    let format = match hap_format {
        Some(hap::HapFormat::Bc1) => gpu::ImgFormat::BC1_RGBA,
        Some(hap::HapFormat::Bc3) => gpu::ImgFormat::BC3_RGBA,
        None => gpu::ImgFormat::R8G8B8A8,
    };

    log::info!(
        "{}x{}, {} frames, codec_tag={}, format={:?}, tsv image: {:?}",
        info.width,
        info.height,
        num_frames,
        info.codec_tag,
        format,
        name
    );

    // Request BC compression features if needed.
    let features = if hap_format.is_some() {
        wgpu::Features::TEXTURE_COMPRESSION_BC
            | if use_array {
                wgpu::Features::empty()
            } else {
                wgpu::Features::TEXTURE_COMPRESSION_BC_SLICED_3D
            }
    } else {
        wgpu::Features::empty()
    };

    let img_type = if use_array {
        gpu::ImgType::D2
    } else {
        gpu::ImgType::D3
    };

    let instance = gpu::create_instance();
    let adapter = gpu::create_adapter(&instance, None);
    let (device, queue) = gpu::create_device_with_features(&adapter, features);

    let mut client = gpu::create_tsv_client(&device);
    client
        .init_image(
            &name,
            info.width,
            info.height,
            num_frames,
            format,
            img_type,
            true,
        )
        .expect("init tsv image");

    let fence = gpu::create_fence(&client);

    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some("video_buffer"),
        size: wgpu::Extent3d {
            width: info.width,
            height: info.height,
            depth_or_array_layers: num_frames,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: gpu::img_type_to_wgpu(img_type, num_frames),
        format: gpu::img_format_to_wgpu(format),
        usage: wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::COPY_SRC,
        view_formats: &[],
    });

    if let Some(hap_fmt) = hap_format {
        upload_hap_frames(&device, &queue, &texture, ff_args, &info, hap_fmt);
    } else {
        upload_raw_frames(&device, &queue, &texture, ff_args, &info);
    }

    // Submit all pending writes and wait for GPU
    queue.submit([]);
    device.poll(wgpu::PollType::wait_indefinitely()).unwrap();

    // Share the complete texture via TSV
    let raw = unsafe { gpu::raw_image(&texture) };
    if let Err(e) = client.send_image(
        &name,
        raw,
        vk::ImageLayout::TRANSFER_DST_OPTIMAL,
        vk::ImageLayout::TRANSFER_DST_OPTIMAL,
        fence,
    ) {
        log::error!("send_image error: {e}");
        std::process::exit(1);
    }

    log::info!("all {} frames uploaded and shared, waiting...", num_frames);

    // Keep the process alive so the shared image remains available
    loop {
        std::thread::sleep(std::time::Duration::from_secs(3600));
    }
}

fn upload_hap_frames(
    _device: &wgpu::Device,
    queue: &wgpu::Queue,
    texture: &wgpu::Texture,
    ff_args: &[String],
    info: &ffmpeg::VideoInfo,
    hap_fmt: hap::HapFormat,
) {
    // For HAP, we need raw packet data (not decoded pixels).
    // Use ffmpeg-next to read packets directly.
    ffmpeg_next::init().expect("ffmpeg init");

    let (_, input_url) = parse_ff_input_args(ff_args);
    let mut ictx = ffmpeg_next::format::input(input_url).expect("open input");

    let video_stream_index = ictx
        .streams()
        .best(ffmpeg_next::media::Type::Video)
        .expect("no video stream")
        .index();

    let compressed_frame_size = hap_fmt.compressed_size(info.width, info.height);
    let mut decode_buf = vec![0u8; compressed_frame_size + 1024 * 1024]; // extra margin
    let mut frame_index: u32 = 0;
    let num_frames = info.num_frames.unwrap();

    let block_bytes = hap_fmt.block_bytes();
    let blocks_x = info.width.div_ceil(4);

    for (stream, packet) in ictx.packets() {
        if stream.index() != video_stream_index {
            continue;
        }

        let data = packet.data().expect("packet data");
        let (_, decoded_len) = hap::decode(data, &mut decode_buf).expect("HAP decode");

        queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture,
                mip_level: 0,
                origin: wgpu::Origin3d {
                    x: 0,
                    y: 0,
                    z: frame_index,
                },
                aspect: wgpu::TextureAspect::All,
            },
            &decode_buf[..decoded_len],
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                // BC: bytes_per_row = blocks_x * block_bytes
                bytes_per_row: Some(blocks_x * block_bytes),
                // rows_per_image = height in blocks
                rows_per_image: Some(info.height.div_ceil(4)),
            },
            wgpu::Extent3d {
                width: info.width,
                height: info.height,
                depth_or_array_layers: 1,
            },
        );

        frame_index += 1;
        if frame_index.is_multiple_of(10) || frame_index == num_frames {
            log::info!("uploaded frame {}/{}", frame_index, num_frames);
        }
        if frame_index >= num_frames {
            break;
        }
    }

    if frame_index < num_frames {
        log::warn!(
            "expected {} frames but only decoded {frame_index}",
            num_frames
        );
    }
}

fn upload_raw_frames(
    _device: &wgpu::Device,
    queue: &wgpu::Queue,
    texture: &wgpu::Texture,
    ff_args: &[String],
    info: &ffmpeg::VideoInfo,
) {
    let frame_size = (info.width * info.height * 4) as usize;
    let num_frames = info.num_frames.unwrap();

    let mut ffmpeg = Command::new("ffmpeg")
        .args(["-v", "quiet"])
        .args(ff_args)
        .args(["-f", "rawvideo", "-pix_fmt", "rgba", "pipe:1"])
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .spawn()
        .expect("failed to start ffmpeg");

    let stdout = ffmpeg.stdout.take().unwrap();
    let mut reader = std::io::BufReader::new(stdout);
    let mut frame_buf = vec![0u8; frame_size];
    let mut frame_index: u32 = 0;

    while reader.read_exact(&mut frame_buf).is_ok() {
        queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture,
                mip_level: 0,
                origin: wgpu::Origin3d {
                    x: 0,
                    y: 0,
                    z: frame_index,
                },
                aspect: wgpu::TextureAspect::All,
            },
            &frame_buf,
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(info.width * 4),
                rows_per_image: Some(info.height),
            },
            wgpu::Extent3d {
                width: info.width,
                height: info.height,
                depth_or_array_layers: 1,
            },
        );

        frame_index += 1;
        if frame_index.is_multiple_of(10) || frame_index == num_frames {
            log::info!("uploaded frame {}/{}", frame_index, num_frames);
        }
        if frame_index >= num_frames {
            break;
        }
    }

    // Drop the reader/stdout to close the pipe — this signals ffmpeg to exit
    // (important when --frames truncates before ffmpeg finishes)
    drop(reader);
    ffmpeg.wait().ok();

    if frame_index < num_frames {
        log::warn!(
            "expected {} frames but only decoded {frame_index}",
            num_frames
        );
    }
}

/// Parse ffmpeg-style input args to extract -f <format> and the input URL.
/// Returns (Option<format>, input_url).
fn parse_ff_input_args(ff_args: &[String]) -> (Option<String>, &str) {
    let mut format = None;
    let mut input_url = None;
    let mut i = 0;

    while i < ff_args.len() {
        match ff_args[i].as_str() {
            "-f" => {
                format = Some(ff_args[i + 1].clone());
                i += 2;
            }
            "-i" => {
                input_url = Some(ff_args[i + 1].as_str());
                i += 2;
            }
            _ => {
                // Last positional arg is input URL if no -i was given
                if i == ff_args.len() - 1 && input_url.is_none() {
                    input_url = Some(ff_args[i].as_str());
                }
                i += 1;
            }
        }
    }

    (
        format,
        input_url.expect("no input file specified in ffmpeg args"),
    )
}