use std::io::Read;
use std::process::{Command, Stdio};
use std::thread;
use std::time::{Duration, Instant};

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

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-stream [--name NAME] -- <ffmpeg input args>");
            std::process::exit(1);
        }
    };

    let mut name = "tsv-video-stream".to_string();

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

    let info = ffmpeg::probe_video(ff_args, None);
    log::info!(
        "{}x{} @ {:.2}fps, tsv image: {name}",
        info.width,
        info.height,
        info.fps
    );

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

    let mut client = gpu::create_tsv_client(&device);
    client
        .init_image(
            &name,
            info.width,
            info.height,
            1,
            gpu::ImgFormat::R8G8B8A8,
            gpu::ImgType::D2,
            true,
        )
        .expect("init tsv image");

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

    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some("video_frame"),
        size: wgpu::Extent3d {
            width: info.width,
            height: info.height,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: wgpu::TextureFormat::Rgba8Unorm,
        usage: wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::COPY_SRC,
        view_formats: &[],
    });

    let frame_size = (info.width * info.height * 4) as usize;
    let frame_time = Duration::from_secs_f64(1.0 / info.fps);

    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];

    loop {
        let frame_start = Instant::now();

        if reader.read_exact(&mut frame_buf).is_err() {
            break;
        }

        queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            &frame_buf,
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(info.width * 4),
                rows_per_image: None,
            },
            wgpu::Extent3d {
                width: info.width,
                height: info.height,
                depth_or_array_layers: 1,
            },
        );
        // flush the pending write_texture command and wait for GPU completion
        queue.submit([]);
        device.poll(wgpu::PollType::wait_indefinitely()).unwrap();

        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::warn!("send_image error: {e}");
        }

        let elapsed = frame_start.elapsed();
        if elapsed < frame_time {
            thread::sleep(frame_time - elapsed);
        }
    }

    ffmpeg.wait().ok();
}