use bone_types::{SketchDimensionId, SketchEntityId, SketchRelationId};
use slotmap::Key;
use std::collections::HashMap;
use std::collections::hash_map::Entry;

use crate::camera::{ViewportExtent, ViewportPx};
use crate::gpu::{Gpu, PICK_BYTES_PER_PIXEL};
use crate::{RenderError, Result};

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum EntityKindTag {
    Point = 0,
    Line = 1,
    Arc = 2,
    Circle = 3,
    Relation = 4,
    Dimension = 5,
}

impl EntityKindTag {
    #[must_use]
    pub const fn bits(self) -> u8 {
        self as u8
    }

    #[must_use]
    pub const fn from_bits(bits: u8) -> Option<Self> {
        match bits {
            0 => Some(Self::Point),
            1 => Some(Self::Line),
            2 => Some(Self::Arc),
            3 => Some(Self::Circle),
            4 => Some(Self::Relation),
            5 => Some(Self::Dimension),
            _ => None,
        }
    }
}

impl core::fmt::Display for EntityKindTag {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let name = match self {
            Self::Point => "point",
            Self::Line => "line",
            Self::Arc => "arc",
            Self::Circle => "circle",
            Self::Relation => "relation",
            Self::Dimension => "dimension",
        };
        f.write_str(name)
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum PickedItem {
    Point(SketchEntityId),
    Line(SketchEntityId),
    Arc(SketchEntityId),
    Circle(SketchEntityId),
    Relation(SketchRelationId),
    Dimension(SketchDimensionId),
}

impl PickedItem {
    #[must_use]
    pub const fn tag(self) -> EntityKindTag {
        match self {
            Self::Point(_) => EntityKindTag::Point,
            Self::Line(_) => EntityKindTag::Line,
            Self::Arc(_) => EntityKindTag::Arc,
            Self::Circle(_) => EntityKindTag::Circle,
            Self::Relation(_) => EntityKindTag::Relation,
            Self::Dimension(_) => EntityKindTag::Dimension,
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum PickIdError {
    #[error("slotmap slot {slot:#x} exceeds {limit:#x}, the 28-bit PickId index field")]
    SlotIndexOverflow { slot: u32, limit: u32 },
    #[error("two entries collapse to packed slot index {0:#x} in the PickIndex")]
    SlotIndexCollision(u32),
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct PickId(u32);

impl PickId {
    pub const TAG_SHIFT: u32 = 28;
    pub const INDEX_MASK: u32 = 0x0FFF_FFFF;
    pub const TAG_MASK: u32 = 0xF000_0000;
    pub const NONE: Self = Self(u32::MAX);

    #[must_use]
    pub const fn from_raw(bits: u32) -> Self {
        Self(bits)
    }

    #[must_use]
    pub const fn raw(self) -> u32 {
        self.0
    }

    #[must_use]
    #[allow(
        clippy::cast_possible_truncation,
        reason = "right-shift by TAG_SHIFT (28) leaves the high nibble in the low 4 bits; value is in 0..16 so u8 fits"
    )]
    pub const fn tag(self) -> Option<EntityKindTag> {
        EntityKindTag::from_bits((self.0 >> Self::TAG_SHIFT) as u8)
    }

    #[must_use]
    pub const fn index(self) -> u32 {
        self.0 & Self::INDEX_MASK
    }

    pub fn point(id: SketchEntityId) -> Result<Self, PickIdError> {
        Self::from_keyed(EntityKindTag::Point, id)
    }

    pub fn line(id: SketchEntityId) -> Result<Self, PickIdError> {
        Self::from_keyed(EntityKindTag::Line, id)
    }

    pub fn arc(id: SketchEntityId) -> Result<Self, PickIdError> {
        Self::from_keyed(EntityKindTag::Arc, id)
    }

    pub fn circle(id: SketchEntityId) -> Result<Self, PickIdError> {
        Self::from_keyed(EntityKindTag::Circle, id)
    }

    pub fn relation(id: SketchRelationId) -> Result<Self, PickIdError> {
        Self::from_keyed(EntityKindTag::Relation, id)
    }

    pub fn dimension(id: SketchDimensionId) -> Result<Self, PickIdError> {
        Self::from_keyed(EntityKindTag::Dimension, id)
    }

    #[must_use]
    pub fn unpack(self, index: &PickIndex) -> Option<PickedItem> {
        let tag = self.tag()?;
        let slot = self.index();
        match tag {
            EntityKindTag::Point => index.lookup_entity(slot, tag).map(PickedItem::Point),
            EntityKindTag::Line => index.lookup_entity(slot, tag).map(PickedItem::Line),
            EntityKindTag::Arc => index.lookup_entity(slot, tag).map(PickedItem::Arc),
            EntityKindTag::Circle => index.lookup_entity(slot, tag).map(PickedItem::Circle),
            EntityKindTag::Relation => index
                .relations
                .get(&slot)
                .copied()
                .map(PickedItem::Relation),
            EntityKindTag::Dimension => index
                .dimensions
                .get(&slot)
                .copied()
                .map(PickedItem::Dimension),
        }
    }

    fn from_keyed<K: Key>(tag: EntityKindTag, id: K) -> Result<Self, PickIdError> {
        slot_index(id).map(|idx| Self::pack(tag, idx))
    }

    const fn pack(tag: EntityKindTag, idx: u32) -> Self {
        debug_assert!(idx <= Self::INDEX_MASK);
        Self(((tag.bits() as u32) << Self::TAG_SHIFT) | idx)
    }
}

const _: () = assert!(PickId::TAG_MASK | PickId::INDEX_MASK == u32::MAX);
const _: () = assert!(PickId::TAG_MASK & PickId::INDEX_MASK == 0);
const _: () = assert!(PickId::TAG_MASK == 0xFu32 << PickId::TAG_SHIFT);
const _: () = assert!(PickId::NONE.tag().is_none());

impl core::fmt::Display for PickId {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self.tag() {
            Some(tag) => write!(f, "{tag}#{}", self.index()),
            None => write!(f, "invalid#{:08x}", self.0),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct PickAperture(u32);

impl PickAperture {
    pub const EXACT: Self = Self(0);
    pub const DEFAULT: Self = Self(5);

    #[must_use]
    pub const fn new(radius_px: u32) -> Self {
        Self(radius_px)
    }

    #[must_use]
    pub const fn radius_px(self) -> u32 {
        self.0
    }
}

impl Default for PickAperture {
    fn default() -> Self {
        Self::DEFAULT
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct PickQuery {
    x: ViewportPx,
    y: ViewportPx,
    aperture: PickAperture,
}

impl PickQuery {
    #[must_use]
    pub const fn new(x: ViewportPx, y: ViewportPx) -> Self {
        Self {
            x,
            y,
            aperture: PickAperture::DEFAULT,
        }
    }

    #[must_use]
    pub const fn with_aperture(self, aperture: PickAperture) -> Self {
        Self { aperture, ..self }
    }

    #[must_use]
    pub const fn x(self) -> ViewportPx {
        self.x
    }

    #[must_use]
    pub const fn y(self) -> ViewportPx {
        self.y
    }

    #[must_use]
    pub const fn aperture(self) -> PickAperture {
        self.aperture
    }
}

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

pub struct Picker<'a> {
    gpu: &'a Gpu,
    pick: &'a wgpu::Texture,
    staging: &'a wgpu::Buffer,
    extent: ViewportExtent,
    index: PickIndex,
}

impl<'a> Picker<'a> {
    pub(crate) fn new(
        gpu: &'a Gpu,
        pick: &'a wgpu::Texture,
        staging: &'a wgpu::Buffer,
        extent: ViewportExtent,
        index: PickIndex,
    ) -> Self {
        Self {
            gpu,
            pick,
            staging,
            extent,
            index,
        }
    }

    pub fn raw_at(&self, query: PickQuery) -> Result<PickId> {
        let width = self.extent.width().value();
        let height = self.extent.height().value();
        if query.x.value() >= width || query.y.value() >= height {
            return Err(RenderError::PickOutOfBounds {
                query,
                extent: self.extent,
            });
        }
        let device = self.gpu.device();
        let queue = self.gpu.queue();
        let padded_bpr = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bone-render:pick-encoder"),
        });
        encoder.copy_texture_to_buffer(
            wgpu::TexelCopyTextureInfo {
                texture: self.pick,
                mip_level: 0,
                origin: wgpu::Origin3d {
                    x: query.x.value(),
                    y: query.y.value(),
                    z: 0,
                },
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::TexelCopyBufferInfo {
                buffer: self.staging,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(padded_bpr),
                    rows_per_image: Some(1),
                },
            },
            wgpu::Extent3d {
                width: 1,
                height: 1,
                depth_or_array_layers: 1,
            },
        );
        queue.submit(Some(encoder.finish()));

        let slice = self.staging.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 raw = {
            let view = slice.get_mapped_range();
            debug_assert!(view.len() >= PICK_BYTES_PER_PIXEL as usize);
            u32::from_le_bytes([view[0], view[1], view[2], view[3]])
        };
        self.staging.unmap();
        Ok(PickId::from_raw(raw))
    }

    pub fn at(&self, query: PickQuery) -> Result<Option<PickedItem>> {
        let radius = query.aperture.radius_px();
        if radius == 0 {
            let pid = self.raw_at(query)?;
            return Ok(pid.unpack(&self.index));
        }
        self.bounds_check(query)?;
        let candidates = self.read_region(query, radius)?;
        Ok(nearest_pick(query, &candidates).and_then(|pid| pid.unpack(&self.index)))
    }

    fn bounds_check(&self, query: PickQuery) -> Result<()> {
        let width = self.extent.width().value();
        let height = self.extent.height().value();
        if query.x.value() >= width || query.y.value() >= height {
            return Err(RenderError::PickOutOfBounds {
                query,
                extent: self.extent,
            });
        }
        Ok(())
    }

    #[allow(
        clippy::cast_possible_wrap,
        clippy::cast_sign_loss,
        reason = "viewport extents and aperture radius fit i32; bounds are clamped non-negative"
    )]
    fn read_region(&self, query: PickQuery, radius: u32) -> Result<Vec<(i32, i32, PickId)>> {
        let r = radius as i32;
        let qx = query.x.value() as i32;
        let qy = query.y.value() as i32;
        let viewport_w = self.extent.width().value() as i32;
        let viewport_h = self.extent.height().value() as i32;
        let x0 = (qx - r).max(0);
        let y0 = (qy - r).max(0);
        let x1 = (qx + r + 1).min(viewport_w);
        let y1 = (qy + r + 1).min(viewport_h);
        let region_w = (x1 - x0) as u32;
        let region_h = (y1 - y0) as u32;
        let raw = self.copy_region_to_host(x0 as u32, y0 as u32, region_w, region_h)?;
        let bpr = padded_row_bytes(region_w);
        Ok(decode_region(x0, y0, region_w, region_h, bpr, &raw))
    }

    fn copy_region_to_host(
        &self,
        origin_x: u32,
        origin_y: u32,
        width: u32,
        height: u32,
    ) -> Result<Vec<u8>> {
        let bpr = padded_row_bytes(width);
        let device = self.gpu.device();
        let queue = self.gpu.queue();
        let staging = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("bone-render:pick-aperture-staging"),
            size: u64::from(bpr) * u64::from(height),
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bone-render:pick-aperture-encoder"),
        });
        encoder.copy_texture_to_buffer(
            wgpu::TexelCopyTextureInfo {
                texture: self.pick,
                mip_level: 0,
                origin: wgpu::Origin3d {
                    x: origin_x,
                    y: origin_y,
                    z: 0,
                },
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::TexelCopyBufferInfo {
                buffer: &staging,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(bpr),
                    rows_per_image: Some(height),
                },
            },
            wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
        );
        queue.submit(Some(encoder.finish()));

        let slice = staging.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 bytes = slice.get_mapped_range().to_vec();
        staging.unmap();
        Ok(bytes)
    }
}

const fn padded_row_bytes(pixels: u32) -> u32 {
    let raw = pixels * PICK_BYTES_PER_PIXEL;
    let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
    raw.div_ceil(align) * align
}

#[allow(
    clippy::cast_possible_wrap,
    reason = "decoded coords are bounded by viewport extents which fit i32"
)]
fn decode_region(
    origin_x: i32,
    origin_y: i32,
    width: u32,
    height: u32,
    bytes_per_row: u32,
    raw: &[u8],
) -> Vec<(i32, i32, PickId)> {
    let bpr = bytes_per_row as usize;
    let bpp = PICK_BYTES_PER_PIXEL as usize;
    assert_eq!(
        raw.len(),
        bpr * height as usize,
        "pick readback length mismatch: got {}, want {} ({} rows of {} bytes)",
        raw.len(),
        bpr * height as usize,
        height,
        bpr,
    );
    (0..height)
        .flat_map(|row| (0..width).map(move |col| (row, col)))
        .filter_map(|(row, col)| {
            let offset = (row as usize) * bpr + (col as usize) * bpp;
            let bytes = raw.get(offset..offset + bpp)?;
            let pid =
                PickId::from_raw(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]));
            Some((origin_x + col as i32, origin_y + row as i32, pid))
        })
        .collect()
}

#[allow(
    clippy::cast_possible_wrap,
    reason = "query coords fit i32 (display pixels)"
)]
#[must_use]
pub(crate) fn nearest_pick(query: PickQuery, candidates: &[(i32, i32, PickId)]) -> Option<PickId> {
    let qx = query.x.value() as i32;
    let qy = query.y.value() as i32;
    candidates
        .iter()
        .copied()
        .filter(|(_, _, pid)| *pid != PickId::NONE)
        .min_by_key(|(px, py, pid)| {
            let dx = px - qx;
            let dy = py - qy;
            let sq_dist = dx * dx + dy * dy;
            let tag_priority = pid.tag().map_or(u8::MAX, EntityKindTag::bits);
            (sq_dist, tag_priority)
        })
        .map(|(_, _, pid)| pid)
}

impl core::fmt::Debug for Picker<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Picker")
            .field("extent", &self.extent)
            .finish_non_exhaustive()
    }
}

#[derive(Debug, Default, Clone)]
pub struct PickIndex {
    entities: HashMap<u32, (SketchEntityId, EntityKindTag)>,
    relations: HashMap<u32, SketchRelationId>,
    dimensions: HashMap<u32, SketchDimensionId>,
}

impl PickIndex {
    pub fn build<EI, RI, DI>(
        entities: EI,
        relations: RI,
        dimensions: DI,
    ) -> Result<Self, PickIdError>
    where
        EI: IntoIterator<Item = (SketchEntityId, EntityKindTag)>,
        RI: IntoIterator<Item = SketchRelationId>,
        DI: IntoIterator<Item = SketchDimensionId>,
    {
        Ok(Self {
            entities: try_collect_keyed(entities, |(id, _)| slot_index(*id))?,
            relations: try_collect_keyed(relations, |id| slot_index(*id))?,
            dimensions: try_collect_keyed(dimensions, |id| slot_index(*id))?,
        })
    }

    fn lookup_entity(&self, slot: u32, expected: EntityKindTag) -> Option<SketchEntityId> {
        self.entities
            .get(&slot)
            .filter(|(_, kind)| *kind == expected)
            .map(|(id, _)| *id)
    }
}

#[allow(
    clippy::cast_possible_truncation,
    reason = "high 32 bits are masked off before the u64-to-u32 cast, so the low 32 bits fit u32 exactly"
)]
fn slot_index<K: Key>(key: K) -> Result<u32, PickIdError> {
    let slot = (key.data().as_ffi() & 0xFFFF_FFFF_u64) as u32;
    if slot > PickId::INDEX_MASK {
        return Err(PickIdError::SlotIndexOverflow {
            slot,
            limit: PickId::INDEX_MASK,
        });
    }
    Ok(slot)
}

fn try_collect_keyed<I, V, F>(iter: I, slot_of: F) -> Result<HashMap<u32, V>, PickIdError>
where
    I: IntoIterator<Item = V>,
    F: Fn(&V) -> Result<u32, PickIdError>,
{
    iter.into_iter().try_fold(HashMap::new(), |mut acc, item| {
        let slot = slot_of(&item)?;
        match acc.entry(slot) {
            Entry::Vacant(v) => {
                v.insert(item);
                Ok(acc)
            }
            Entry::Occupied(_) => Err(PickIdError::SlotIndexCollision(slot)),
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use proptest::prelude::*;
    use slotmap::{KeyData, SlotMap};

    #[derive(Copy, Clone, Debug)]
    enum TagKind {
        Point,
        Line,
        Arc,
        Circle,
        Relation,
        Dimension,
    }

    impl TagKind {
        fn entity_kind(self) -> EntityKindTag {
            match self {
                Self::Point | Self::Relation | Self::Dimension => EntityKindTag::Point,
                Self::Line => EntityKindTag::Line,
                Self::Arc => EntityKindTag::Arc,
                Self::Circle => EntityKindTag::Circle,
            }
        }
    }

    fn arb_tag() -> impl Strategy<Value = TagKind> {
        prop_oneof![
            Just(TagKind::Point),
            Just(TagKind::Line),
            Just(TagKind::Arc),
            Just(TagKind::Circle),
            Just(TagKind::Relation),
            Just(TagKind::Dimension),
        ]
    }

    type TestSlots = (
        SlotMap<SketchEntityId, EntityKindTag>,
        SlotMap<SketchRelationId, ()>,
        SlotMap<SketchDimensionId, ()>,
    );

    fn populate_uniform(
        n: usize,
        entity_kind: EntityKindTag,
    ) -> (
        Vec<SketchEntityId>,
        Vec<SketchRelationId>,
        Vec<SketchDimensionId>,
        TestSlots,
    ) {
        let mut entities = SlotMap::with_key();
        let mut relations = SlotMap::with_key();
        let mut dimensions = SlotMap::with_key();
        let e: Vec<_> = (0..n).map(|_| entities.insert(entity_kind)).collect();
        let r: Vec<_> = (0..n).map(|_| relations.insert(())).collect();
        let d: Vec<_> = (0..n).map(|_| dimensions.insert(())).collect();
        (e, r, d, (entities, relations, dimensions))
    }

    fn ok<T>(result: Result<T, PickIdError>) -> T {
        let Ok(value) = result else {
            panic!("test fixture should fit in PickId field");
        };
        value
    }

    fn build_index(slots: &TestSlots) -> PickIndex {
        ok(PickIndex::build(
            slots.0.iter().map(|(k, kind)| (k, *kind)),
            slots.1.keys(),
            slots.2.keys(),
        ))
    }

    fn build_case(
        tag: TagKind,
        e: &[SketchEntityId],
        r: &[SketchRelationId],
        d: &[SketchDimensionId],
        target: usize,
    ) -> (PickId, PickedItem) {
        match tag {
            TagKind::Point => (ok(PickId::point(e[target])), PickedItem::Point(e[target])),
            TagKind::Line => (ok(PickId::line(e[target])), PickedItem::Line(e[target])),
            TagKind::Arc => (ok(PickId::arc(e[target])), PickedItem::Arc(e[target])),
            TagKind::Circle => (ok(PickId::circle(e[target])), PickedItem::Circle(e[target])),
            TagKind::Relation => (
                ok(PickId::relation(r[target])),
                PickedItem::Relation(r[target]),
            ),
            TagKind::Dimension => (
                ok(PickId::dimension(d[target])),
                PickedItem::Dimension(d[target]),
            ),
        }
    }

    fn forget(slots: &mut TestSlots, item: PickedItem) {
        match item {
            PickedItem::Point(k)
            | PickedItem::Line(k)
            | PickedItem::Arc(k)
            | PickedItem::Circle(k) => {
                slots.0.remove(k);
            }
            PickedItem::Relation(k) => {
                slots.1.remove(k);
            }
            PickedItem::Dimension(k) => {
                slots.2.remove(k);
            }
        }
    }

    #[test]
    fn tag_bits_roundtrip() {
        (0u8..6).for_each(|bits| {
            let Some(tag) = EntityKindTag::from_bits(bits) else {
                panic!("tag {bits} should decode");
            };
            assert_eq!(tag.bits(), bits);
        });
    }

    #[test]
    fn invalid_tag_bits_return_none() {
        (6u8..16).for_each(|bits| {
            assert!(EntityKindTag::from_bits(bits).is_none());
        });
    }

    #[test]
    fn raw_roundtrip() {
        let mut entities: SlotMap<SketchEntityId, ()> = SlotMap::with_key();
        let k = entities.insert(());
        let pid = ok(PickId::point(k));
        assert_eq!(PickId::from_raw(pid.raw()), pid);
    }

    #[test]
    fn tag_and_index_split_cleanly() {
        let mut entities: SlotMap<SketchEntityId, ()> = SlotMap::with_key();
        let k = entities.insert(());
        let pid = ok(PickId::line(k));
        assert_eq!(pid.tag(), Some(EntityKindTag::Line));
        assert_eq!(pid.raw() & PickId::TAG_MASK, 1u32 << PickId::TAG_SHIFT);
        assert_eq!(pid.index(), pid.raw() & PickId::INDEX_MASK);
    }

    #[test]
    fn display_renders_tag_and_index() {
        let mut entities: SlotMap<SketchEntityId, ()> = SlotMap::with_key();
        let k = entities.insert(());
        let pid = ok(PickId::arc(k));
        assert!(format!("{pid}").starts_with("arc#"));
    }

    #[test]
    fn display_renders_invalid_tag() {
        let raw = 0xF000_0000u32;
        let pid = PickId::from_raw(raw);
        assert_eq!(pid.tag(), None);
        assert_eq!(format!("{pid}"), "invalid#f0000000");
    }

    #[test]
    fn max_index_roundtrip() {
        let raw =
            (u32::from(EntityKindTag::Point.bits()) << PickId::TAG_SHIFT) | PickId::INDEX_MASK;
        let pid = PickId::from_raw(raw);
        assert_eq!(pid.tag(), Some(EntityKindTag::Point));
        assert_eq!(pid.index(), PickId::INDEX_MASK);
        assert_eq!(pid.raw(), raw);
    }

    #[test]
    fn slot_overflow_is_rejected_by_constructor() {
        let bad_slot: u32 = PickId::INDEX_MASK + 1;
        let key = SketchEntityId::from(KeyData::from_ffi((1u64 << 32) | u64::from(bad_slot)));
        let result = PickId::point(key);
        assert!(matches!(
            result,
            Err(PickIdError::SlotIndexOverflow { slot, limit })
                if slot == bad_slot && limit == PickId::INDEX_MASK
        ));
    }

    #[test]
    fn slot_overflow_is_rejected_by_pick_index_build() {
        let bad_slot: u32 = PickId::INDEX_MASK + 1;
        let key = SketchEntityId::from(KeyData::from_ffi((1u64 << 32) | u64::from(bad_slot)));
        let result = PickIndex::build(
            [(key, EntityKindTag::Point)],
            std::iter::empty::<SketchRelationId>(),
            std::iter::empty::<SketchDimensionId>(),
        );
        assert!(matches!(result, Err(PickIdError::SlotIndexOverflow { .. })));
    }

    #[test]
    fn none_sentinel_is_invalid() {
        assert_eq!(PickId::NONE.raw(), u32::MAX);
        assert_eq!(PickId::NONE.tag(), None);
        assert_eq!(PickId::NONE.unpack(&PickIndex::default()), None);
    }

    #[test]
    fn collision_in_pick_index_build_is_rejected() {
        let mut entities: SlotMap<SketchEntityId, EntityKindTag> = SlotMap::with_key();
        let k = entities.insert(EntityKindTag::Point);
        let result = PickIndex::build(
            [(k, EntityKindTag::Point), (k, EntityKindTag::Line)],
            std::iter::empty::<SketchRelationId>(),
            std::iter::empty::<SketchDimensionId>(),
        );
        assert!(matches!(result, Err(PickIdError::SlotIndexCollision(_))));
    }

    fn build_entity_index(entities: &SlotMap<SketchEntityId, EntityKindTag>) -> PickIndex {
        ok(PickIndex::build(
            entities.iter().map(|(k, kind)| (k, *kind)),
            std::iter::empty::<SketchRelationId>(),
            std::iter::empty::<SketchDimensionId>(),
        ))
    }

    #[test]
    fn slot_reuse_with_same_kind_returns_new_key() {
        let mut entities: SlotMap<SketchEntityId, EntityKindTag> = SlotMap::with_key();
        let k1 = entities.insert(EntityKindTag::Point);
        let pid = ok(PickId::point(k1));
        entities.remove(k1);
        let k2 = entities.insert(EntityKindTag::Point);
        assert_ne!(k1, k2);
        let index = build_entity_index(&entities);
        assert_eq!(pid.unpack(&index), Some(PickedItem::Point(k2)));
    }

    #[test]
    fn slot_reuse_with_different_kind_returns_none() {
        let mut entities: SlotMap<SketchEntityId, EntityKindTag> = SlotMap::with_key();
        let k1 = entities.insert(EntityKindTag::Point);
        let pid = ok(PickId::point(k1));
        entities.remove(k1);
        let _k2 = entities.insert(EntityKindTag::Line);
        let index = build_entity_index(&entities);
        assert_eq!(pid.unpack(&index), None);
    }

    proptest! {
        #[test]
        fn unpack_roundtrip_across_all_populated(
            tag in arb_tag(),
            n in 1usize..32,
            target in 0usize..32,
        ) {
            let target = target % n;
            let (e, r, d, slots) = populate_uniform(n, tag.entity_kind());
            let (pid, expected) = build_case(tag, &e, &r, &d, target);
            let index = build_index(&slots);
            prop_assert_eq!(pid.unpack(&index), Some(expected));
            prop_assert_eq!(pid.tag(), Some(expected.tag()));
        }

        #[test]
        fn removed_slot_unpacks_to_none(
            tag in arb_tag(),
            n in 1usize..32,
            target in 0usize..32,
        ) {
            let target = target % n;
            let (e, r, d, mut slots) = populate_uniform(n, tag.entity_kind());
            let (pid, item) = build_case(tag, &e, &r, &d, target);
            forget(&mut slots, item);
            let index = build_index(&slots);
            prop_assert!(pid.unpack(&index).is_none());
        }

        #[test]
        fn invalid_tag_bits_unpack_to_none(bits in 6u8..16, idx in 0u32..=PickId::INDEX_MASK) {
            let raw = (u32::from(bits) << PickId::TAG_SHIFT) | (idx & PickId::INDEX_MASK);
            let pid = PickId::from_raw(raw);
            prop_assert!(pid.unpack(&PickIndex::default()).is_none());
        }
    }

    fn pid_for(slot: u32, tag: EntityKindTag) -> PickId {
        debug_assert!(slot <= PickId::INDEX_MASK);
        PickId::from_raw((u32::from(tag.bits()) << PickId::TAG_SHIFT) | slot)
    }

    fn query_at(x: u32, y: u32) -> PickQuery {
        PickQuery::new(ViewportPx::new(x), ViewportPx::new(y))
    }

    #[test]
    fn aperture_default_is_five_pixels() {
        assert_eq!(PickAperture::default().radius_px(), 5);
        assert_eq!(PickAperture::EXACT.radius_px(), 0);
    }

    #[test]
    fn pick_query_carries_default_aperture() {
        let q = query_at(10, 20);
        assert_eq!(q.aperture(), PickAperture::DEFAULT);
        let exact = q.with_aperture(PickAperture::EXACT);
        assert_eq!(exact.aperture(), PickAperture::EXACT);
        assert_eq!((exact.x(), exact.y()), (q.x(), q.y()));
    }

    #[test]
    fn nearest_pick_returns_none_when_all_empty() {
        let q = query_at(5, 5);
        let region = vec![
            (4, 5, PickId::NONE),
            (5, 5, PickId::NONE),
            (6, 5, PickId::NONE),
        ];
        assert_eq!(nearest_pick(q, &region), None);
    }

    #[test]
    fn nearest_pick_picks_center_when_center_hits() {
        let q = query_at(5, 5);
        let center_pid = pid_for(7, EntityKindTag::Line);
        let neighbour_pid = pid_for(11, EntityKindTag::Line);
        let region = vec![
            (3, 5, neighbour_pid),
            (5, 5, center_pid),
            (8, 5, neighbour_pid),
        ];
        assert_eq!(nearest_pick(q, &region), Some(center_pid));
    }

    #[test]
    fn nearest_pick_resolves_by_squared_distance() {
        let q = query_at(10, 10);
        let near_pid = pid_for(1, EntityKindTag::Line);
        let far_pid = pid_for(2, EntityKindTag::Line);
        let region = vec![(13, 14, far_pid), (12, 11, near_pid)];
        assert_eq!(nearest_pick(q, &region), Some(near_pid));
    }

    #[test]
    fn nearest_pick_breaks_ties_by_iteration_order() {
        let q = query_at(0, 0);
        let first_pid = pid_for(1, EntityKindTag::Line);
        let second_pid = pid_for(2, EntityKindTag::Line);
        let region = vec![(3, 4, first_pid), (4, 3, second_pid)];
        assert_eq!(nearest_pick(q, &region), Some(first_pid));
    }

    #[test]
    fn nearest_pick_skips_none_and_picks_only_real() {
        let q = query_at(20, 20);
        let real = pid_for(7, EntityKindTag::Arc);
        let region = vec![
            (20, 20, PickId::NONE),
            (22, 22, real),
            (25, 25, PickId::NONE),
        ];
        assert_eq!(nearest_pick(q, &region), Some(real));
    }

    #[test]
    fn padded_row_bytes_aligns_to_wgpu_alignment() {
        assert_eq!(
            padded_row_bytes(1),
            wgpu::COPY_BYTES_PER_ROW_ALIGNMENT,
            "1 px row pads to alignment"
        );
        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
        let pixels = align / PICK_BYTES_PER_PIXEL;
        assert_eq!(
            padded_row_bytes(pixels),
            align,
            "exactly one alignment-worth of pixels stays at one alignment unit"
        );
        assert_eq!(
            padded_row_bytes(pixels + 1),
            align * 2,
            "one extra pixel triggers a second alignment unit"
        );
    }

    #[test]
    fn decode_region_shifts_indices_by_origin() {
        let bpr_u32 = padded_row_bytes(2);
        let bpr = bpr_u32 as usize;
        let mut raw = vec![0u8; bpr * 2];
        let pid_a = pid_for(1, EntityKindTag::Line);
        let pid_b = pid_for(2, EntityKindTag::Arc);
        raw[0..4].copy_from_slice(&pid_a.raw().to_le_bytes());
        raw[bpr + 4..bpr + 8].copy_from_slice(&pid_b.raw().to_le_bytes());
        let decoded = decode_region(10, 20, 2, 2, bpr_u32, &raw);
        assert_eq!(decoded.len(), 4);
        assert_eq!(decoded[0], (10, 20, pid_a));
        assert_eq!(decoded[3], (11, 21, pid_b));
    }
}
