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

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

/// Options that ffmpeg supports but ffprobe does not.
const FFMPEG_ONLY_OPTIONS: &[&str] = &["-stream_loop", "-t"];

fn filter_probe_args(ff_args: &[String]) -> Vec<String> {
    let mut out = Vec::new();
    let mut skip_next = false;
    for arg in ff_args {
        if skip_next {
            skip_next = false;
            continue;
        }
        if FFMPEG_ONLY_OPTIONS.contains(&arg.as_str()) {
            skip_next = true;
            continue;
        }
        out.push(arg.clone());
    }
    out
}

struct VideoInfo {
    width: u32,
    height: u32,
    num_frames: u32,
    codec_tag: String,
}

fn probe_video(ff_args: &[String], max_frames: Option<u32>) -> VideoInfo {
    let probe_args = filter_probe_args(ff_args);
    let output = Command::new("ffprobe")
        .args(["-v", "error", "-select_streams", "v:0"])
        .args([
            "-show_entries",
            "stream=width,height,nb_frames,codec_tag_string",
        ])
        .args(["-of", "csv=p=0"])
        .args(&probe_args)
        .output()
        .expect("failed to run ffprobe");

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        panic!("ffprobe failed: {stderr}");
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let line = stdout.trim();
    let parts: Vec<&str> = line.split(',').collect();
    if parts.len() < 4 {
        panic!("unexpected ffprobe output: {line}");
    }

    let num_frames = parts[3].parse().unwrap_or_else(|_| {
        max_frames.unwrap_or_else(|| {
            panic!(
                "could not determine frame count (nb_frames={}).\n\
                 Use --frames N or use tsv-video-stream instead.",
                parts[3]
            )
        })
    });

    VideoInfo {
        codec_tag: parts[0].to_string(),
        width: parts[1].parse().expect("parse width"),
        height: parts[2].parse().expect("parse height"),
        num_frames,
    }
}

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 = probe_video(ff_args, max_frames);
    if let Some(max) = max_frames {
        info.num_frames = info.num_frames.min(max);
    }
    let hap_format = is_hap_codec(&info.codec_tag);

    let (texture_format, tsv_img_format) = match hap_format {
        Some(hap::HapFormat::Bc1) => (wgpu::TextureFormat::Bc1RgbaUnorm, gpu::ImgFormat::BC1_RGBA),
        Some(hap::HapFormat::Bc3) => (wgpu::TextureFormat::Bc3RgbaUnorm, gpu::ImgFormat::BC3_RGBA),
        None => (
            wgpu::TextureFormat::Rgba8UnormSrgb,
            gpu::ImgFormat::R8G8B8A8,
        ),
    };

    log::info!(
        "{}x{}, {} frames, codec_tag={}, format={:?}, tsv image: {}",
        info.width,
        info.height,
        info.num_frames,
        info.codec_tag,
        texture_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 (tsv_img_type, wgpu_dimension) = if use_array {
        (gpu::ImgType::D2, wgpu::TextureDimension::D2)
    } else {
        (gpu::ImgType::D3, wgpu::TextureDimension::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,
            info.num_frames,
            tsv_img_format,
            tsv_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: info.num_frames,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu_dimension,
        format: texture_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...",
        info.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: &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 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 == info.num_frames {
            log::info!("uploaded frame {}/{}", frame_index, info.num_frames);
        }
        if frame_index >= info.num_frames {
            break;
        }
    }

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

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

    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 == info.num_frames {
            log::info!("uploaded frame {}/{}", frame_index, info.num_frames);
        }
        if frame_index >= info.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 < info.num_frames {
        log::warn!(
            "expected {} frames but only decoded {frame_index}",
            info.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"),
    )
}