use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fmt;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread::ThreadId;

use crate::changes::MAX_PENDING_CHANGES;
use crate::{ChangeBatch, ChangeKey, ObjectChange};
use crate::{EndorObject, ObjectHeader, ObjectPayload, Value, VmError};

static NEXT_HEAP_ID: AtomicU64 = AtomicU64::new(1);

/// Checked, process-local reference. Copying it never retains an object.
/// The generation prevents a released handle from aliasing a reused slot.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ObjectRef {
    heap_id: u64,
    index: usize,
    generation: u64,
}

impl fmt::Display for ObjectRef {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}:{}", self.heap_id, self.index, self.generation)
    }
}

/// A non-owning reference. Only the originating store can upgrade it; slot
/// reuse never makes a dead weak reference point at a different object.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct WeakObjectRef(pub(crate) ObjectRef);

#[derive(Debug)]
struct Slot {
    generation: u64,
    object: Option<Box<EndorObject>>,
}

#[derive(Debug)]
struct ReadOnly {
    protected: BTreeSet<ObjectRef>,
    reads: Mutex<BTreeSet<ObjectRef>>,
    regions: Mutex<Vec<BTreeMap<ObjectRef, bool>>>,
    dependencies: Mutex<BTreeMap<ReadKey, ReadDependency>>,
    dependency_regions: Mutex<Vec<BTreeMap<ReadKey, ReadDependency>>>,
}

#[derive(Debug)]
pub(crate) struct Heap {
    id: u64,
    slots: Vec<Slot>,
    free: Vec<usize>,
    live: usize,
    changes: BTreeMap<(ObjectRef, ChangeKey), u64>,
    changes_overflow: bool,
    read_only: Vec<ReadOnly>,
    validation_reads: Mutex<HashMap<ThreadId, usize>>,
}

/// Borrowed read identity; content_read distinguishes ordinary reads from
/// annotation validation alone. Neither form retains the object.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ReadAccess {
    pub object: ObjectRef,
    pub content_read: bool,
}

/// Borrowed observations; Object handles never retain the observed graph.
/// Named values are immutable comparison snapshots, not managed owning edges.
#[derive(Debug, Clone, PartialEq)]
pub enum ReadDependency {
    Object {
        object: ObjectRef,
        version: u64,
    },
    DictionaryKind {
        object: ObjectRef,
    },
    Name {
        object: ObjectRef,
        name: String,
        value: Option<crate::Value>,
    },
}
impl ReadDependency {
    pub fn object(&self) -> ObjectRef {
        match self {
            Self::Object { object, .. }
            | Self::DictionaryKind { object }
            | Self::Name { object, .. } => *object,
        }
    }
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
enum ReadKey {
    Object(ObjectRef),
    DictionaryKind(ObjectRef),
    Name(ObjectRef, String),
}
#[derive(Debug)]
pub struct TrackedReads {
    pub accesses: Vec<ReadAccess>,
    pub dependencies: Vec<ReadDependency>,
}
#[derive(Debug)]
pub struct ReadSet {
    pub objects: Vec<ObjectRef>,
    pub dependencies: Vec<ReadDependency>,
}

struct ValidationScope<'a> {
    heap: &'a Heap,
    thread: ThreadId,
}
impl Drop for ValidationScope<'_> {
    fn drop(&mut self) {
        let mut active = self
            .heap
            .validation_reads
            .lock()
            .expect("read tracker poisoned");
        let depth = active
            .get_mut(&self.thread)
            .expect("active validation scope");
        *depth -= 1;
        if *depth == 0 {
            active.remove(&self.thread);
        }
    }
}

impl Heap {
    pub(crate) fn new() -> Result<Self, VmError> {
        let id = NEXT_HEAP_ID
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
            .map_err(|_| VmError::IdentityExhausted)?;
        Ok(Self {
            id,
            slots: Vec::new(),
            free: Vec::new(),
            live: 0,
            changes: BTreeMap::new(),
            changes_overflow: false,
            read_only: vec![],
            validation_reads: Mutex::new(HashMap::new()),
        })
    }

    pub(crate) fn live_references(&self) -> impl Iterator<Item = ObjectRef> + '_ {
        self.slots.iter().enumerate().filter_map(|(index, slot)| {
            slot.object.as_ref().map(|_| ObjectRef {
                heap_id: self.id,
                index,
                generation: slot.generation,
            })
        })
    }

    pub(crate) fn begin_read_only(&mut self) -> usize {
        let protected = self
            .slots
            .iter()
            .enumerate()
            .filter_map(|(index, slot)| {
                slot.object.as_ref().map(|_| ObjectRef {
                    heap_id: self.id,
                    index,
                    generation: slot.generation,
                })
            })
            .collect();
        self.read_only.push(ReadOnly {
            protected,
            reads: Mutex::new(BTreeSet::new()),
            regions: Mutex::new(Vec::new()),
            dependencies: Mutex::new(BTreeMap::new()),
            dependency_regions: Mutex::new(Vec::new()),
        });
        self.read_only.len()
    }

    pub(crate) fn end_read_only(&mut self, scope: usize) -> Result<Vec<ObjectRef>, VmError> {
        Ok(self.end_read_only_precise(scope)?.objects)
    }
    pub(crate) fn end_read_only_precise(&mut self, scope: usize) -> Result<ReadSet, VmError> {
        if scope == 0 || scope != self.read_only.len() {
            return Err(VmError::InvalidArgument("read-only scope nesting mismatch"));
        }
        let scope = self.read_only.pop().unwrap();
        Ok(ReadSet {
            objects: scope
                .reads
                .into_inner()
                .expect("read tracker poisoned")
                .into_iter()
                .collect(),
            dependencies: scope
                .dependencies
                .into_inner()
                .expect("read tracker poisoned")
                .into_values()
                .collect(),
        })
    }

    pub(crate) fn begin_read_tracking(&self) -> Result<usize, VmError> {
        let scope = self.read_only.last().ok_or(VmError::InvalidArgument(
            "read tracking requires a read-only scope",
        ))?;
        let mut regions = scope.regions.lock().expect("read tracker poisoned");
        regions.push(BTreeMap::new());
        scope
            .dependency_regions
            .lock()
            .expect("read tracker poisoned")
            .push(BTreeMap::new());
        Ok(regions.len())
    }

    pub(crate) fn end_read_tracking(&self, token: usize) -> Result<Vec<ObjectRef>, VmError> {
        Ok(self
            .end_read_tracking_details(token)?
            .into_iter()
            .map(|read| read.object)
            .collect())
    }

    pub(crate) fn end_read_tracking_details(
        &self,
        token: usize,
    ) -> Result<Vec<ReadAccess>, VmError> {
        Ok(self.end_read_tracking_precise(token)?.accesses)
    }
    pub(crate) fn end_read_tracking_precise(&self, token: usize) -> Result<TrackedReads, VmError> {
        let scope = self.read_only.last().ok_or(VmError::InvalidArgument(
            "read tracking requires a read-only scope",
        ))?;
        let mut regions = scope.regions.lock().expect("read tracker poisoned");
        let mut precise = scope
            .dependency_regions
            .lock()
            .expect("read tracker poisoned");
        if token == 0 || token != regions.len() || token != precise.len() {
            return Err(VmError::InvalidArgument("read tracking nesting mismatch"));
        }
        Ok(TrackedReads {
            accesses: regions
                .pop()
                .unwrap()
                .into_iter()
                .map(|(object, content_read)| ReadAccess {
                    object,
                    content_read,
                })
                .collect(),
            dependencies: precise.pop().unwrap().into_values().collect(),
        })
    }

    pub(crate) fn is_read_only(&self) -> bool {
        !self.read_only.is_empty()
    }

    pub(crate) fn with_validation_reads<T>(
        &self,
        compute: impl FnOnce() -> Result<T, VmError>,
    ) -> Result<T, VmError> {
        if self.read_only.is_empty() {
            return compute();
        }
        let thread = std::thread::current().id();
        {
            let mut active = self.validation_reads.lock().expect("read tracker poisoned");
            let depth = active.entry(thread).or_default();
            *depth = depth
                .checked_add(1)
                .ok_or(VmError::ExecutionLimit("validation scopes"))?;
        }
        let _scope = ValidationScope { heap: self, thread };
        compute()
    }

    pub(crate) fn record_read(&self, reference: ObjectRef, version: u64) {
        if self.read_only.is_empty() {
            return;
        }
        self.record_observation(reference, ReadKey::Object(reference), || {
            ReadDependency::Object {
                object: reference,
                version,
            }
        });
    }
    pub(crate) fn record_dictionary_kind(&self, reference: ObjectRef) {
        self.record_observation(reference, ReadKey::DictionaryKind(reference), || {
            ReadDependency::DictionaryKind { object: reference }
        });
    }
    pub(crate) fn record_name(
        &self,
        reference: ObjectRef,
        name: &str,
        value: &Option<crate::Value>,
    ) {
        if self.read_only.is_empty() {
            return;
        }
        self.record_observation(reference, ReadKey::Name(reference, name.into()), || {
            ReadDependency::Name {
                object: reference,
                name: name.into(),
                value: value.clone(),
            }
        });
    }
    fn record_observation(
        &self,
        reference: ObjectRef,
        key: ReadKey,
        build: impl Fn() -> ReadDependency,
    ) {
        if self.read_only.is_empty() {
            return;
        }
        let content = !self
            .validation_reads
            .lock()
            .expect("read tracker poisoned")
            .contains_key(&std::thread::current().id());
        for scope in &self.read_only {
            if scope.protected.contains(&reference) {
                scope
                    .dependencies
                    .lock()
                    .expect("read tracker poisoned")
                    .entry(key.clone())
                    .or_insert_with(&build);
                for region in scope
                    .dependency_regions
                    .lock()
                    .expect("read tracker poisoned")
                    .iter_mut()
                {
                    region.entry(key.clone()).or_insert_with(&build);
                }
                scope
                    .reads
                    .lock()
                    .expect("read tracker poisoned")
                    .insert(reference);
                for region in scope
                    .regions
                    .lock()
                    .expect("read tracker poisoned")
                    .iter_mut()
                {
                    region
                        .entry(reference)
                        .and_modify(|ordinary| *ordinary |= content)
                        .or_insert(content);
                }
            }
        }
    }

    pub(crate) fn allocate(
        &mut self,
        payload: Box<dyn ObjectPayload>,
    ) -> Result<ObjectRef, VmError> {
        let mut edges = Vec::new();
        payload.trace(&mut |reference| edges.push(reference));
        let increments = self.prepare_increments(&edges)?;
        let index = self.free.pop().unwrap_or(self.slots.len());
        if index == self.slots.len() {
            self.slots.push(Slot {
                generation: 0,
                object: None,
            });
        }
        let reference = ObjectRef {
            heap_id: self.id,
            index,
            generation: self.slots[index].generation,
        };
        let object = Box::new(EndorObject {
            header: ObjectHeader {
                identity: reference,
                strong_refs: 1,
                external_refs: 1,
                change_version: 0,
                type_object: None,
            },
            payload,
            attributes: BTreeMap::new(),
            payload_edges: edges,
        });
        self.apply_increments(increments);
        self.slots[index].object = Some(object);
        self.live += 1;
        Ok(reference)
    }

    pub(crate) fn get(&self, reference: ObjectRef) -> Result<&EndorObject, VmError> {
        self.validate_owner(reference)?;
        self.slots
            .get(reference.index)
            .filter(|slot| slot.generation == reference.generation)
            .and_then(|slot| slot.object.as_deref())
            .ok_or(VmError::InvalidHandle)
    }

    pub(crate) fn attach_type(
        &mut self,
        reference: ObjectRef,
        type_object: ObjectRef,
    ) -> Result<(), VmError> {
        if self.get(reference)?.header.type_object == Some(type_object) {
            return Ok(());
        }
        if self.get(reference)?.header.type_object.is_some() {
            return Err(VmError::InvalidArgument("object already has a type"));
        }
        if self.get(reference)?.attributes.contains_key("type") {
            return Err(VmError::InvalidArgument(
                "type attribute conflicts with managed type",
            ));
        }
        let increments = self.prepare_increments(&[type_object])?;
        self.prepare_change(reference)?;
        self.apply_increments(increments);
        self.get_mut(reference)?.header.type_object = Some(type_object);
        self.record_change(reference, ChangeKey::Contents);
        Ok(())
    }

    /// Transactional copy-on-update for built-in containers. The EndorObject
    /// box and identity stay stable; only its payload and edge snapshot change.
    /// All new edges are validated/retained before any old edge is released.
    pub(crate) fn replace_payload(
        &mut self,
        reference: ObjectRef,
        payload: impl ObjectPayload,
    ) -> Result<(), VmError> {
        self.get(reference)?;
        let mut edges = Vec::new();
        payload.trace(&mut |target| edges.push(target));
        let increments = self.prepare_increments(&edges)?;
        self.prepare_change(reference)?;
        self.apply_increments(increments);
        let object = self.get_mut(reference)?;
        let old_edges = std::mem::replace(&mut object.payload_edges, edges);
        let old_payload = std::mem::replace(&mut object.payload, Box::new(payload));
        self.record_change(reference, ChangeKey::Contents);
        self.release_edges(old_edges);
        drop(old_payload);
        Ok(())
    }

    fn get_mut(&mut self, reference: ObjectRef) -> Result<&mut EndorObject, VmError> {
        self.get(reference)?;
        Ok(self.slots[reference.index].object.as_deref_mut().unwrap())
    }

    fn validate_owner(&self, reference: ObjectRef) -> Result<(), VmError> {
        if reference.heap_id != self.id {
            return Err(VmError::ForeignHandle);
        }
        Ok(())
    }

    fn prepare_increments(
        &self,
        references: &[ObjectRef],
    ) -> Result<BTreeMap<ObjectRef, usize>, VmError> {
        let mut counts = BTreeMap::<ObjectRef, usize>::new();
        for reference in references {
            self.get(*reference)?;
            let count = counts.entry(*reference).or_default();
            *count = count.checked_add(1).ok_or(VmError::RefCountOverflow)?;
        }
        for (reference, count) in &counts {
            self.get(*reference)?
                .header
                .strong_refs
                .checked_add(*count)
                .ok_or(VmError::RefCountOverflow)?;
        }
        Ok(counts)
    }

    fn apply_increments(&mut self, increments: BTreeMap<ObjectRef, usize>) {
        for (reference, count) in increments {
            self.get_mut(reference).unwrap().header.strong_refs += count;
        }
    }

    pub(crate) fn retain(&mut self, reference: ObjectRef) -> Result<(), VmError> {
        let header = &mut self.get_mut(reference)?.header;
        let strong = header
            .strong_refs
            .checked_add(1)
            .ok_or(VmError::RefCountOverflow)?;
        let external = header
            .external_refs
            .checked_add(1)
            .ok_or(VmError::RefCountOverflow)?;
        header.strong_refs = strong;
        header.external_refs = external;
        Ok(())
    }

    pub(crate) fn release(&mut self, reference: ObjectRef) -> Result<(), VmError> {
        let header = &mut self.get_mut(reference)?.header;
        if header.external_refs == 0 {
            return Err(VmError::NoOwnedReference);
        }
        header.external_refs -= 1;
        header.strong_refs -= 1;
        let pending = if header.strong_refs == 0 {
            vec![reference]
        } else {
            Vec::new()
        };
        self.drain_zero_counts(pending);
        Ok(())
    }

    pub(crate) fn set_attr(
        &mut self,
        reference: ObjectRef,
        name: String,
        value: Value,
    ) -> Result<(), VmError> {
        self.check_attributes(reference)?;
        let edges: Vec<_> = match &value {
            Value::Object(target) => vec![*target],
            _ => Vec::new(),
        };
        let increments = self.prepare_increments(&edges)?;
        self.prepare_change(reference)?;
        self.apply_increments(increments);
        let old = self
            .get_mut(reference)?
            .attributes
            .insert(name.clone(), value);
        self.record_change(reference, ChangeKey::Attribute(name));
        if let Some(Value::Object(target)) = old {
            self.release_edges(vec![target]);
        }
        Ok(())
    }

    pub(crate) fn get_attr(&self, reference: ObjectRef, name: &str) -> Result<Value, VmError> {
        self.check_attributes(reference)?;
        self.get(reference)?
            .attributes
            .get(name)
            .cloned()
            .ok_or_else(|| VmError::AttributeNotFound(name.into()))
    }

    pub(crate) fn del_attr(&mut self, reference: ObjectRef, name: &str) -> Result<(), VmError> {
        self.check_attributes(reference)?;
        self.prepare_change(reference)?;
        let old = self
            .get_mut(reference)?
            .attributes
            .remove(name)
            .ok_or_else(|| VmError::AttributeNotFound(name.into()))?;
        self.record_change(reference, ChangeKey::Attribute(name.into()));
        if let Value::Object(target) = old {
            self.release_edges(vec![target]);
        }
        Ok(())
    }

    pub(crate) fn prepare_change(&self, reference: ObjectRef) -> Result<(), VmError> {
        self.get(reference)?;
        if self
            .read_only
            .iter()
            .any(|scope| scope.protected.contains(&reference))
        {
            return Err(VmError::UnsupportedOperation(
                "mutation of existing object during read-only call",
            ));
        }
        self.get(reference)?
            .header
            .change_version
            .checked_add(1)
            .ok_or(VmError::ExecutionLimit("object change version"))?;
        Ok(())
    }

    pub(crate) fn record_change(&mut self, reference: ObjectRef, key: ChangeKey) {
        let object = self.get_mut(reference).expect("validated change target");
        object.header.change_version += 1;
        let version = object.header.change_version;
        // Fresh result construction has no existing observers. Keep versions,
        // but do not notify application dependencies about derived allocations.
        if self.is_read_only() {
            return;
        }
        if self.changes_overflow {
            return;
        }
        let entry = (reference, key);
        if self.changes.len() >= MAX_PENDING_CHANGES && !self.changes.contains_key(&entry) {
            self.changes.clear();
            self.changes_overflow = true;
        } else {
            self.changes.insert(entry, version);
        }
    }

    pub(crate) fn take_changes(&mut self) -> ChangeBatch {
        ChangeBatch {
            changes: std::mem::take(&mut self.changes)
                .into_iter()
                .map(|((object, key), version)| ObjectChange {
                    object,
                    key,
                    version,
                })
                .collect(),
            rescan_required: std::mem::take(&mut self.changes_overflow),
        }
    }

    fn check_attributes(&self, reference: ObjectRef) -> Result<(), VmError> {
        if !self.get(reference)?.payload.allows_attributes() {
            return Err(VmError::UnsupportedOperation("attributes"));
        }
        Ok(())
    }

    fn release_edges(&mut self, edges: Vec<ObjectRef>) {
        let mut pending = Vec::new();
        for target in edges {
            let header = &mut self.get_mut(target).unwrap().header;
            header.strong_refs -= 1;
            if header.strong_refs == 0 {
                pending.push(target);
            }
        }
        self.drain_zero_counts(pending);
    }

    /// Removes the object before running Drop. Reuse increments its generation;
    /// an exhausted generation permanently retires that slot.
    fn take(&mut self, reference: ObjectRef) -> Box<EndorObject> {
        let slot = &mut self.slots[reference.index];
        let object = slot.object.take().unwrap();
        if let Some(next) = slot.generation.checked_add(1) {
            slot.generation = next;
            self.free.push(reference.index);
        }
        self.live -= 1;
        object
    }

    fn drain_zero_counts(&mut self, mut pending: Vec<ObjectRef>) {
        while let Some(reference) = pending.pop() {
            let object = self.take(reference);
            for target in object.edges() {
                let header = &mut self.get_mut(target).unwrap().header;
                header.strong_refs -= 1;
                if header.strong_refs == 0 {
                    pending.push(target);
                }
            }
            drop(object);
        }
    }

    /// Full, explicit cycle collection. External references are roots; every
    /// managed attribute and declared payload edge participates in tracing.
    pub(crate) fn collect_cycles(&mut self) -> usize {
        let mut marked = vec![false; self.slots.len()];
        let mut pending = Vec::new();
        for slot in &self.slots {
            if let Some(object) = &slot.object
                && object.header.external_refs > 0
            {
                pending.push(object.header.identity);
            }
        }
        while let Some(reference) = pending.pop() {
            if marked[reference.index] {
                continue;
            }
            marked[reference.index] = true;
            pending.extend(self.get(reference).unwrap().edges());
        }
        let garbage: Vec<_> = self
            .slots
            .iter()
            .enumerate()
            .filter(|(index, _)| !marked[*index])
            .filter_map(|(_, slot)| slot.object.as_ref().map(|object| object.header.identity))
            .collect();
        let objects: Vec<_> = garbage
            .iter()
            .map(|reference| self.take(*reference))
            .collect();
        // All garbage slots are detached before touching edges or running Drop.
        for object in &objects {
            for target in object.edges() {
                if let Ok(survivor) = self.get_mut(target) {
                    survivor.header.strong_refs -= 1;
                }
            }
        }
        let collected = objects.len();
        drop(objects);
        collected
    }

    pub(crate) fn len(&self) -> usize {
        self.live
    }
}

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

    #[test]
    fn invalid_slots_and_generations_are_rejected() {
        let mut heap = Heap::new().unwrap();
        let reference = heap.allocate(Box::new(Record)).unwrap();
        let invalid = ObjectRef {
            index: usize::MAX,
            ..reference
        };
        let stale = ObjectRef {
            generation: 1,
            ..reference
        };
        assert!(matches!(heap.get(invalid), Err(VmError::InvalidHandle)));
        assert!(matches!(heap.get_mut(stale), Err(VmError::InvalidHandle)));
    }

    #[test]
    fn reused_slots_never_alias_and_exhausted_generations_retire() {
        let mut heap = Heap::new().unwrap();
        let first = heap.allocate(Box::new(Record)).unwrap();
        heap.release(first).unwrap();
        let second = heap.allocate(Box::new(Record)).unwrap();
        assert_eq!(first.index, second.index);
        assert_ne!(first.generation, second.generation);
        assert!(matches!(heap.get(first), Err(VmError::InvalidHandle)));
        heap.slots[second.index].generation = u64::MAX;
        let last = ObjectRef {
            generation: u64::MAX,
            ..second
        };
        heap.slots[last.index]
            .object
            .as_mut()
            .unwrap()
            .header
            .identity = last;
        heap.release(last).unwrap();
        let next = heap.allocate(Box::new(Record)).unwrap();
        assert_ne!(next.index, last.index);
    }

    #[test]
    fn failed_reference_increments_leave_counts_and_attributes_unchanged() {
        let mut heap = Heap::new().unwrap();
        let owner = heap.allocate(Box::new(Record)).unwrap();
        let target = heap.allocate(Box::new(Record)).unwrap();
        heap.get_mut(target).unwrap().header.strong_refs = usize::MAX;
        assert_eq!(heap.retain(target), Err(VmError::RefCountOverflow));
        assert_eq!(heap.get(target).unwrap().header.external_refs, 1);
        assert_eq!(
            heap.set_attr(owner, "edge".into(), Value::Object(target)),
            Err(VmError::RefCountOverflow)
        );
        assert!(heap.get(owner).unwrap().attributes.is_empty());
    }
}