use crate::camera::{ViewportExtent, ViewportPx};
use crate::pick::{PickId, PickIndex, Picker};
use crate::snapshot::{SnapshotFrame, Style};
use crate::{RenderError, Result};

pub(crate) const COLOR_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8Unorm;
pub(crate) const PICK_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::R32Uint;
pub(crate) const BYTES_PER_PIXEL: u32 = 4;
pub(crate) const PICK_BYTES_PER_PIXEL: u32 = 4;
const MSAA_SAMPLE_COUNT: u32 = 4;

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum AdapterPolicy {
    Platform,
    Software,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct BackendTag(wgpu::Backend);

impl BackendTag {
    #[cfg(test)]
    pub(crate) const fn from_backend(backend: wgpu::Backend) -> Self {
        Self(backend)
    }

    #[must_use]
    pub fn backend(self) -> wgpu::Backend {
        self.0
    }
}

impl core::fmt::Display for BackendTag {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{:?}", self.0)
    }
}

#[derive(Clone, Debug)]
pub struct Capabilities {
    adapter_limits: wgpu::Limits,
    backend: BackendTag,
    adapter_name: String,
    msaa_x4: bool,
}

impl Capabilities {
    pub(crate) fn probe(adapter: &wgpu::Adapter) -> Self {
        let info = adapter.get_info();
        let color = adapter.get_texture_format_features(COLOR_FORMAT).flags;
        let depth = adapter
            .get_texture_format_features(crate::pipelines::solid::DEPTH_FORMAT)
            .flags;
        Self {
            adapter_limits: adapter.limits(),
            backend: BackendTag(info.backend),
            adapter_name: info.name,
            msaa_x4: msaa_x4_supported(color, depth),
        }
    }

    #[must_use]
    pub fn adapter_limits(&self) -> &wgpu::Limits {
        &self.adapter_limits
    }

    #[must_use]
    pub fn backend(&self) -> BackendTag {
        self.backend
    }

    #[must_use]
    pub fn adapter_name(&self) -> &str {
        &self.adapter_name
    }

    #[must_use]
    pub fn supports_msaa_x4(&self) -> bool {
        self.msaa_x4
    }
}

fn msaa_x4_supported(
    color: wgpu::TextureFormatFeatureFlags,
    depth: wgpu::TextureFormatFeatureFlags,
) -> bool {
    color.sample_count_supported(MSAA_SAMPLE_COUNT)
        && depth.sample_count_supported(MSAA_SAMPLE_COUNT)
}

pub struct Gpu {
    device: wgpu::Device,
    queue: wgpu::Queue,
    capabilities: Capabilities,
}

impl Gpu {
    pub(crate) fn from_parts(
        device: wgpu::Device,
        queue: wgpu::Queue,
        capabilities: Capabilities,
    ) -> Self {
        Self {
            device,
            queue,
            capabilities,
        }
    }

    #[must_use]
    pub fn device(&self) -> &wgpu::Device {
        &self.device
    }

    #[must_use]
    pub fn queue(&self) -> &wgpu::Queue {
        &self.queue
    }

    #[must_use]
    pub fn capabilities(&self) -> &Capabilities {
        &self.capabilities
    }
}

pub struct OffscreenContext {
    gpu: Gpu,
    color: wgpu::Texture,
    pick: wgpu::Texture,
    depth: wgpu::Texture,
    pick_staging: wgpu::Buffer,
    extent: ViewportExtent,
}

fn require_nonzero(extent: ViewportExtent) -> Result<()> {
    if extent.width().value() == 0 || extent.height().value() == 0 {
        return Err(RenderError::ZeroExtent);
    }
    Ok(())
}

fn offscreen_color_texture(device: &wgpu::Device, extent: ViewportExtent) -> wgpu::Texture {
    device.create_texture(&wgpu::TextureDescriptor {
        label: Some("bone-render:offscreen-color"),
        size: texture_size(extent),
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: COLOR_FORMAT,
        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
        view_formats: &[],
    })
}

fn offscreen_pick_texture(device: &wgpu::Device, extent: ViewportExtent) -> wgpu::Texture {
    device.create_texture(&wgpu::TextureDescriptor {
        label: Some("bone-render:offscreen-pick"),
        size: texture_size(extent),
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: PICK_FORMAT,
        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
        view_formats: &[],
    })
}

impl OffscreenContext {
    pub async fn new(extent: ViewportExtent, policy: AdapterPolicy) -> Result<Self> {
        require_nonzero(extent)?;
        let instance =
            wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle_from_env());
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::LowPower,
                force_fallback_adapter: matches!(policy, AdapterPolicy::Software),
                compatible_surface: None,
            })
            .await?;
        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                label: Some("bone-render:offscreen"),
                required_features: wgpu::Features::empty(),
                required_limits: wgpu::Limits::downlevel_defaults(),
                experimental_features: wgpu::ExperimentalFeatures::default(),
                memory_hints: wgpu::MemoryHints::default(),
                trace: wgpu::Trace::Off,
            })
            .await?;
        let capabilities = Capabilities::probe(&adapter);
        let color = offscreen_color_texture(&device, extent);
        let pick = offscreen_pick_texture(&device, extent);
        let depth = crate::depth_texture(&device, extent);
        let pick_staging = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("bone-render:pick-readback"),
            size: u64::from(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT),
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        Ok(Self {
            gpu: Gpu {
                device,
                queue,
                capabilities,
            },
            color,
            pick,
            depth,
            pick_staging,
            extent,
        })
    }

    pub fn resize(&mut self, extent: ViewportExtent) -> Result<()> {
        require_nonzero(extent)?;
        self.color = offscreen_color_texture(&self.gpu.device, extent);
        self.pick = offscreen_pick_texture(&self.gpu.device, extent);
        self.depth = crate::depth_texture(&self.gpu.device, extent);
        self.extent = extent;
        Ok(())
    }

    pub fn render_passes<F>(&self, build_passes: F)
    where
        F: FnOnce(
            &mut wgpu::CommandEncoder,
            &wgpu::TextureView,
            &wgpu::TextureView,
            &wgpu::TextureView,
        ),
    {
        let mut encoder = self
            .gpu
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("bone-render:offscreen-frame"),
            });
        let color_view = self.color_view();
        let pick_view = self.pick_view();
        let depth_view = self
            .depth
            .create_view(&wgpu::TextureViewDescriptor::default());
        build_passes(&mut encoder, &color_view, &pick_view, &depth_view);
        self.gpu.queue.submit(Some(encoder.finish()));
    }

    pub fn capture(&self) -> Result<SnapshotFrame> {
        self.render(|_, _, _| {})
    }

    #[must_use]
    pub fn gpu(&self) -> &Gpu {
        &self.gpu
    }

    #[must_use]
    pub fn extent(&self) -> ViewportExtent {
        self.extent
    }

    #[must_use]
    pub const fn color_format(&self) -> wgpu::TextureFormat {
        COLOR_FORMAT
    }

    #[must_use]
    pub fn color_view(&self) -> wgpu::TextureView {
        self.color
            .create_view(&wgpu::TextureViewDescriptor::default())
    }

    #[must_use]
    pub fn pick_view(&self) -> wgpu::TextureView {
        self.pick
            .create_view(&wgpu::TextureViewDescriptor::default())
    }

    #[must_use]
    pub fn picker(&self, index: PickIndex) -> Picker<'_> {
        Picker::new(
            &self.gpu,
            &self.pick,
            &self.pick_staging,
            self.extent,
            index,
        )
    }

    pub fn render_clear(&self, style: &Style) -> Result<SnapshotFrame> {
        self.render(|encoder, color_view, pick_view| {
            let _pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("bone-render:clear"),
                color_attachments: &[
                    Some(wgpu::RenderPassColorAttachment {
                        view: color_view,
                        resolve_target: None,
                        depth_slice: None,
                        ops: wgpu::Operations {
                            load: wgpu::LoadOp::Clear(style.background().into()),
                            store: wgpu::StoreOp::Store,
                        },
                    }),
                    Some(wgpu::RenderPassColorAttachment {
                        view: pick_view,
                        resolve_target: None,
                        depth_slice: None,
                        ops: wgpu::Operations {
                            load: wgpu::LoadOp::Clear(pick_clear_color()),
                            store: wgpu::StoreOp::Store,
                        },
                    }),
                ],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
                multiview_mask: None,
            });
        })
    }

    pub(crate) fn render<F>(&self, mut build_passes: F) -> Result<SnapshotFrame>
    where
        F: FnMut(&mut wgpu::CommandEncoder, &wgpu::TextureView, &wgpu::TextureView),
    {
        let mut encoder = self
            .gpu
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("bone-render:encoder"),
            });
        let color_view = self.color_view();
        let pick_view = self.pick_view();
        build_passes(&mut encoder, &color_view, &pick_view);
        let padded_bpr = padded_bytes_per_row(self.extent.width());
        let staging = self.gpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("bone-render:readback"),
            size: u64::from(padded_bpr) * u64::from(self.extent.height().value()),
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        encoder.copy_texture_to_buffer(
            wgpu::TexelCopyTextureInfo {
                texture: &self.color,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::TexelCopyBufferInfo {
                buffer: &staging,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(padded_bpr),
                    rows_per_image: Some(self.extent.height().value()),
                },
            },
            texture_size(self.extent),
        );
        self.gpu.queue.submit(Some(encoder.finish()));
        let rgba = read_staging(&self.gpu.device, &staging, self.extent, padded_bpr)?;
        Ok(SnapshotFrame::new(
            self.extent,
            rgba,
            self.gpu.capabilities.backend(),
        ))
    }
}

pub(crate) fn clear_pick_attachment(
    encoder: &mut wgpu::CommandEncoder,
    pick_view: &wgpu::TextureView,
) {
    let _pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
        label: Some("bone-render:pick-clear"),
        color_attachments: &[Some(wgpu::RenderPassColorAttachment {
            view: pick_view,
            resolve_target: None,
            depth_slice: None,
            ops: wgpu::Operations {
                load: wgpu::LoadOp::Clear(pick_clear_color()),
                store: wgpu::StoreOp::Store,
            },
        })],
        depth_stencil_attachment: None,
        timestamp_writes: None,
        occlusion_query_set: None,
        multiview_mask: None,
    });
}

pub(crate) fn pick_clear_color() -> wgpu::Color {
    wgpu::Color {
        r: f64::from(PickId::NONE.raw()),
        g: 0.0,
        b: 0.0,
        a: 0.0,
    }
}

fn texture_size(extent: ViewportExtent) -> wgpu::Extent3d {
    wgpu::Extent3d {
        width: extent.width().value(),
        height: extent.height().value(),
        depth_or_array_layers: 1,
    }
}

pub(crate) fn padded_bytes_per_row(width: ViewportPx) -> u32 {
    let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
    let w = width.value();
    debug_assert!(
        w <= u32::MAX / BYTES_PER_PIXEL,
        "viewport width too large for padded row calc"
    );
    let raw = w * BYTES_PER_PIXEL;
    raw.div_ceil(align) * align
}

fn read_staging(
    device: &wgpu::Device,
    buffer: &wgpu::Buffer,
    extent: ViewportExtent,
    padded_bpr: u32,
) -> Result<Vec<u8>> {
    let slice = buffer.slice(..);
    let (tx, rx) =
        std::sync::mpsc::sync_channel::<core::result::Result<(), wgpu::BufferAsyncError>>(1);
    slice.map_async(wgpu::MapMode::Read, move |res| {
        let _ = tx.send(res);
    });
    device
        .poll(wgpu::PollType::wait_indefinitely())
        .map_err(RenderError::Poll)?;
    match rx.try_recv() {
        Ok(Ok(())) => {}
        Ok(Err(e)) => return Err(RenderError::Map(e)),
        Err(_) => return Err(RenderError::MapMissing),
    }
    let row_bytes = extent.width().value() * BYTES_PER_PIXEL;
    let rgba: Vec<u8> = {
        let view = slice.get_mapped_range();
        (0..extent.height().value())
            .flat_map(|y| {
                let start = (y * padded_bpr) as usize;
                let end = start + row_bytes as usize;
                view[start..end].iter().copied()
            })
            .collect()
    };
    buffer.unmap();
    Ok(rgba)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn msaa_x4_requires_both_color_and_depth_support() {
        let yes = wgpu::TextureFormatFeatureFlags::MULTISAMPLE_X4;
        let no = wgpu::TextureFormatFeatureFlags::empty();
        assert!(msaa_x4_supported(yes, yes));
        assert!(!msaa_x4_supported(yes, no));
        assert!(!msaa_x4_supported(no, yes));
        assert!(!msaa_x4_supported(no, no));
    }
}
