use crate::heap::Heap;
use crate::{
    ChangeBatch, ChangeKey, CompareOp, EndorObject, ObjectPayload, ObjectRef, Record, Value,
    VmError, WeakObjectRef,
};

/// Object storage independent of the VM evaluator. Construction creates an empty
/// heap; dropping it releases all objects and their host resources. Mutable
/// access is limited to operations that validate and account for owned edges.
#[derive(Debug)]
pub struct ObjectStore {
    pub(crate) heap: Heap,
    function_type: Option<WeakObjectRef>,
    pub(crate) field_type_checks: bool,
    pub(crate) active_attribute_checks: std::sync::Mutex<
        std::collections::HashMap<std::thread::ThreadId, std::collections::BTreeSet<ObjectRef>>,
    >,
    pub(crate) active_trait_checks: std::sync::Mutex<
        std::collections::HashMap<
            std::thread::ThreadId,
            std::collections::BTreeSet<(ObjectRef, ObjectRef)>,
        >,
    >,
}

impl ObjectStore {
    /// Opt-in declared-field and method substitution checks, shared with the VM.
    pub fn set_field_type_checks(&mut self, enabled: bool) -> Result<(), VmError> {
        if self.is_read_only() {
            return Err(VmError::UnsupportedOperation(
                "configure field type checks during read-only call",
            ));
        }
        self.field_type_checks = enabled;
        Ok(())
    }

    /// Existing objects are immutable; fresh result graphs may be constructed.
    /// Scopes must end in LIFO order. Read handles are borrowed, never GC roots.
    pub fn begin_read_only(&mut self) -> usize {
        self.heap.begin_read_only()
    }
    pub fn end_read_only(&mut self, scope: usize) -> Result<Vec<ObjectRef>, VmError> {
        self.heap.end_read_only(scope)
    }
    pub fn end_read_only_precise(&mut self, scope: usize) -> Result<crate::ReadSet, VmError> {
        self.heap.end_read_only_precise(scope)
    }
    pub fn read_dependency_valid(&self, dependency: &crate::ReadDependency) -> bool {
        match dependency {
            crate::ReadDependency::Object { object, version } => {
                self.change_version(*object).ok() == Some(*version)
            }
            crate::ReadDependency::DictionaryKind { object } => {
                self.dictionary_kind(*object).is_ok()
            }
            crate::ReadDependency::Name {
                object,
                name,
                value,
            } => self
                .dict_get_name(*object, name)
                .is_ok_and(|current| current == *value),
        }
    }
    pub fn end_read_tracking_precise(&self, token: usize) -> Result<crate::TrackedReads, VmError> {
        self.heap.end_read_tracking_precise(token)
    }
    /// Track one subcall inside the active read-only scope. Regions nest in
    /// LIFO order and do not clear the enclosing call's read dependencies.
    /// Handles returned by end_read_tracking are borrowed, never owning roots.
    pub fn begin_read_tracking(&self) -> Result<usize, VmError> {
        self.heap.begin_read_tracking()
    }
    pub fn end_read_tracking(&self, token: usize) -> Result<Vec<ObjectRef>, VmError> {
        self.heap.end_read_tracking(token)
    }

    /// Like end_read_tracking, retaining the distinction between program reads
    /// and VM annotation validation. Ordinary reads always dominate.
    pub fn end_read_tracking_details(
        &self,
        token: usize,
    ) -> Result<Vec<crate::ReadAccess>, VmError> {
        self.heap.end_read_tracking_details(token)
    }

    pub fn is_read_only(&self) -> bool {
        self.heap.is_read_only()
    }

    pub fn downgrade(&self, reference: ObjectRef) -> Result<WeakObjectRef, VmError> {
        self.object(reference)?;
        Ok(WeakObjectRef(reference))
    }

    /// Returns a new owned root, or None when the original object was released.
    /// A foreign store is an error, not an expired reference.
    pub fn upgrade(&mut self, weak: WeakObjectRef) -> Result<Option<ObjectRef>, VmError> {
        match self.object(weak.0) {
            Ok(_) => {
                self.retain(weak.0)?;
                Ok(Some(weak.0))
            }
            Err(VmError::InvalidHandle) => Ok(None),
            Err(error) => Err(error),
        }
    }

    pub fn take_changes(&mut self) -> ChangeBatch {
        self.heap.take_changes()
    }

    pub fn change_version(&self, reference: ObjectRef) -> Result<u64, VmError> {
        Ok(self.object(reference)?.header().change_version())
    }

    /// Replace a host payload transactionally, including all owned GC edges.
    pub fn replace_payload(
        &mut self,
        reference: ObjectRef,
        payload: impl ObjectPayload,
    ) -> Result<(), VmError> {
        if self
            .object(reference)?
            .payload()
            .as_any()
            .is::<crate::EndorTypeObject>()
        {
            return Err(VmError::UnsupportedOperation(
                "replace immutable type definition",
            ));
        }
        if self
            .object(reference)?
            .payload()
            .as_any()
            .is::<crate::EndorTraitObject>()
        {
            return Err(VmError::UnsupportedOperation(
                "replace immutable trait definition",
            ));
        }
        self.validate_builtin(&payload)?;
        self.heap.replace_payload(reference, payload)
    }

    pub fn assign_from(&mut self, target: ObjectRef, source: ObjectRef) -> Result<bool, VmError> {
        self.object(target)?;
        self.object(source)?;
        if let Some(cell) = self
            .object(target)?
            .payload()
            .assignment_cell(source, self)?
        {
            self.heap.prepare_change(target)?;
            // Replacing the Cell edge may otherwise reclaim a target which is
            // reachable only through the edge being removed.
            self.retain(target)?;
            let result = self.cell_set(cell, Value::Object(source));
            if result.is_ok() {
                self.heap.record_change(target, ChangeKey::Contents);
            }
            self.release(target)?;
            return result.map(|()| true);
        }
        if self.is_read_only() && self.object(target)?.payload().reference_assignment_only() {
            return Ok(false);
        }
        self.heap.prepare_change(target)?;
        let changed = self
            .object(target)?
            .payload()
            .assign_from(self.object(source)?.payload())?;
        if changed {
            self.heap.record_change(target, ChangeKey::Contents);
        }
        Ok(changed)
    }

    pub fn new() -> Result<Self, VmError> {
        Ok(Self {
            heap: Heap::new()?,
            function_type: None,
            field_type_checks: false,
            active_attribute_checks: std::sync::Mutex::new(std::collections::HashMap::new()),
            active_trait_checks: std::sync::Mutex::new(std::collections::HashMap::new()),
        })
    }

    pub fn object_count(&self) -> usize {
        self.heap.len()
    }

    /// Returns one external strong reference. Copies of the returned handle
    /// are borrowed; use retain/release to manage additional owned references.
    pub fn allocate(&mut self, payload: impl ObjectPayload) -> Result<ObjectRef, VmError> {
        self.validate_builtin(&payload)?;
        let function = payload.is_function();
        let object = self.heap.allocate(Box::new(payload))?;
        if function && let Some(weak) = self.function_type {
            let result = (|| {
                let definition = self.upgrade(weak)?.ok_or(VmError::InvalidArgument(
                    "configured Function type has expired",
                ))?;
                let result = self.attach_type(object, definition);
                self.release(definition)?;
                result
            })();
            if let Err(error) = result {
                self.release(object)?;
                return Err(error);
            }
        }
        Ok(object)
    }

    pub fn set_function_type(&mut self, definition: ObjectRef) -> Result<(), VmError> {
        self.payload::<crate::EndorTypeObject>(definition, "Type")?;
        if self.is_read_only() {
            return Err(VmError::UnsupportedOperation(
                "configure Function type during read-only call",
            ));
        }
        let weak = self.downgrade(definition)?;
        if self.function_type.is_some_and(|old| old != weak) {
            return Err(VmError::InvalidArgument(
                "Function type is already configured",
            ));
        }
        self.function_type = Some(weak);
        Ok(())
    }

    pub fn new_object(&mut self) -> Result<ObjectRef, VmError> {
        self.allocate(Record)
    }

    pub fn object(&self, reference: ObjectRef) -> Result<&EndorObject, VmError> {
        let object = self.heap.get(reference)?;
        self.heap
            .record_read(reference, object.header().change_version());
        Ok(object)
    }

    /// Borrowed storage path: depend on Contents for intermediate objects and
    /// Attribute(name) on the final storage object. Every handle is validated.
    pub fn attribute_path(&self, reference: ObjectRef) -> Result<Vec<ObjectRef>, VmError> {
        let mut path = Vec::new();
        let mut seen = std::collections::BTreeSet::new();
        let mut current = reference;
        loop {
            if !seen.insert(current) {
                return Err(VmError::InvalidArgument("cyclic attribute storage"));
            }
            path.push(current);
            match self.object(current)?.payload().attribute_target(self)? {
                Some(target) => current = target,
                None => return Ok(path),
            }
        }
    }

    pub fn retain(&mut self, reference: ObjectRef) -> Result<(), VmError> {
        self.heap.retain(reference)
    }

    pub fn release(&mut self, reference: ObjectRef) -> Result<(), VmError> {
        self.heap.release(reference)
    }

    pub fn ref_count(&self, reference: ObjectRef) -> Result<usize, VmError> {
        Ok(self.heap.get(reference)?.header().ref_count())
    }

    pub fn is_same(&self, left: ObjectRef, right: ObjectRef) -> Result<bool, VmError> {
        self.heap.get(left)?;
        self.heap.get(right)?;
        Ok(left == right)
    }

    pub fn repr(&self, reference: ObjectRef) -> Result<String, VmError> {
        self.heap.get(reference)?.payload().repr(reference)
    }

    pub fn str(&self, reference: ObjectRef) -> Result<String, VmError> {
        self.heap.get(reference)?.payload().str(reference)
    }

    pub fn is_true(&self, reference: ObjectRef) -> Result<bool, VmError> {
        self.heap.get(reference)?.payload().is_true()
    }

    pub fn hash(&self, reference: ObjectRef) -> Result<u64, VmError> {
        self.heap.get(reference)?.payload().hash(reference)
    }

    pub fn compare(
        &self,
        left: ObjectRef,
        right: ObjectRef,
        operation: CompareOp,
    ) -> Result<bool, VmError> {
        let left_payload = self.heap.get(left)?.payload();
        let right_payload = self.heap.get(right)?.payload();
        if let Some(result) = left_payload.compare(right_payload, operation)? {
            return Ok(result);
        }
        if let Some(result) = right_payload.compare(left_payload, operation.reflected())? {
            return Ok(result);
        }
        match operation {
            CompareOp::Eq => Ok(left == right),
            CompareOp::Ne => Ok(left != right),
            _ => Err(VmError::UnsupportedOperation("ordering")),
        }
    }

    /// Reads a borrowed value. An Object result needs retain() if the host
    /// wants it to survive removal of the attribute that currently owns it.
    pub fn get_attr(&self, reference: ObjectRef, name: &str) -> Result<Value, VmError> {
        self.lookup_attr(reference, name)
            .map(|attribute| match attribute {
                crate::types::AttributeLookup::Value(value) => value,
                crate::types::AttributeLookup::Method { function, .. } => Value::Object(function),
            })
    }

    /// Object-valued attributes own internal strong references independently
    /// of the caller's external references. Failure leaves the field unchanged.
    pub fn set_attr(
        &mut self,
        reference: ObjectRef,
        name: impl Into<String>,
        value: Value,
    ) -> Result<(), VmError> {
        let name = name.into();
        if name == "type" && self.type_of(reference)?.is_some() {
            return Err(VmError::UnsupportedOperation("write object type"));
        }
        if self.field_type_checks {
            self.check_field_write(reference, &name, &value)?;
        }
        for object in self.attribute_path(reference)? {
            self.check_typed_method_write(object, &name)?;
        }
        let storage = *self
            .attribute_path(reference)?
            .last()
            .expect("nonempty attribute path");
        self.heap.set_attr(storage, name, value)
    }

    pub fn del_attr(&mut self, reference: ObjectRef, name: &str) -> Result<(), VmError> {
        if self.field_type_checks && self.declared_field(reference, name)?.is_some() {
            return Err(VmError::UnsupportedOperation("delete declared class field"));
        }
        if name == "type" && self.type_of(reference)?.is_some() {
            return Err(VmError::UnsupportedOperation("delete object type"));
        }
        let storage = *self
            .attribute_path(reference)?
            .last()
            .expect("nonempty attribute path");
        self.heap.del_attr(storage, name)
    }

    /// Validate a borrowed object handle without taking an owned reference.
    pub fn validate_value(&self, value: &Value) -> Result<(), VmError> {
        if let Value::Object(reference) = value {
            self.heap.get(*reference)?;
        }
        Ok(())
    }

    /// Explicit full collection. Execution slots use owned roots, so values
    /// on live execution stacks participate in the same root accounting.
    pub fn collect_cycles(&mut self) -> usize {
        self.heap.collect_cycles()
    }
}