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)]
    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 PickQuery {
    x: ViewportPx,
    y: ViewportPx,
}

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

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

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

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 pid = self.raw_at(query)?;
        Ok(pid.unpack(&self.index))
    }
}

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)]
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());
        }
    }
}
