// Fusabi VM Garbage Collection
// Implements mark-and-sweep garbage collection for managing heap-allocated objects

use crate::closure::{Closure, Upvalue};
use crate::value::Value;
use std::collections::HashMap;

/// Trace trait for marking reachable objects during garbage collection
pub trait Trace {
    /// Mark all objects reachable from this object
    fn trace(&self, tracer: &mut Tracer);
}

/// Color states for tri-color marking algorithm
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Color {
    /// White - not yet visited, candidate for collection
    White,
    /// Grey - visited but children not yet scanned
    Grey,
    /// Black - visited and all children scanned
    Black,
}

/// Header attached to all GC-managed objects
#[derive(Debug, Clone)]
pub struct GcHeader {
    /// Current color of the object in the marking phase
    pub color: Color,
    /// Unique identifier for this object
    pub id: usize,
}

impl GcHeader {
    pub fn new(id: usize) -> Self {
        GcHeader {
            color: Color::White,
            id,
        }
    }

    pub fn mark(&mut self) {
        self.color = Color::Grey;
    }

    pub fn is_marked(&self) -> bool {
        self.color != Color::White
    }

    pub fn reset(&mut self) {
        self.color = Color::White;
    }
}

/// Tracer used during the mark phase to track reachable objects
pub struct Tracer {
    /// Grey set - objects marked but not yet scanned
    pub grey_set: Vec<usize>,
    /// Set of all marked object IDs
    pub marked: std::collections::HashSet<usize>,
}

impl Tracer {
    pub fn new() -> Self {
        Tracer {
            grey_set: Vec::new(),
            marked: std::collections::HashSet::new(),
        }
    }

    /// Mark an object as reachable
    pub fn mark(&mut self, id: usize) {
        if !self.marked.contains(&id) {
            self.marked.insert(id);
            self.grey_set.push(id);
        }
    }

    /// Check if there are more objects to scan
    pub fn has_grey(&self) -> bool {
        !self.grey_set.is_empty()
    }

    /// Pop the next grey object to scan
    pub fn pop_grey(&mut self) -> Option<usize> {
        self.grey_set.pop()
    }
}

impl Default for Tracer {
    fn default() -> Self {
        Self::new()
    }
}

/// Wrapper for GC-managed objects
#[derive(Debug, Clone)]
pub struct GcObject {
    pub header: GcHeader,
    pub value: Value,
}

impl GcObject {
    pub fn new(id: usize, value: Value) -> Self {
        GcObject {
            header: GcHeader::new(id),
            value,
        }
    }
}

/// Garbage-collected heap that manages all allocated objects
#[derive(Debug)]
pub struct GcHeap {
    /// All objects currently allocated
    objects: HashMap<usize, GcObject>,
    /// Next object ID to assign
    next_id: usize,
    /// Total bytes allocated (approximation)
    bytes_allocated: usize,
    /// Threshold for triggering collection
    next_gc: usize,
    /// Collection statistics
    pub stats: GcStats,
}

/// Statistics about garbage collection performance
#[derive(Debug, Clone, Default)]
pub struct GcStats {
    /// Total number of collections performed
    pub collections: usize,
    /// Total objects collected
    pub objects_collected: usize,
    /// Total bytes collected
    pub bytes_collected: usize,
}

impl GcHeap {
    /// Create a new GC heap with default threshold (1MB)
    pub fn new() -> Self {
        Self::with_threshold(1024 * 1024)
    }

    /// Create a new GC heap with a custom threshold
    pub fn with_threshold(threshold: usize) -> Self {
        GcHeap {
            objects: HashMap::new(),
            next_id: 0,
            bytes_allocated: 0,
            next_gc: threshold,
            stats: GcStats::default(),
        }
    }

    /// Allocate a new object on the heap
    pub fn allocate(&mut self, value: Value) -> usize {
        let id = self.next_id;
        self.next_id += 1;

        let size = estimate_value_size(&value);
        self.bytes_allocated += size;

        let obj = GcObject::new(id, value);
        self.objects.insert(id, obj);

        id
    }

    /// Get a reference to an object by ID
    pub fn get(&self, id: usize) -> Option<&Value> {
        self.objects.get(&id).map(|obj| &obj.value)
    }

    /// Get the number of objects currently allocated
    pub fn object_count(&self) -> usize {
        self.objects.len()
    }

    /// Get the approximate bytes allocated
    pub fn bytes_allocated(&self) -> usize {
        self.bytes_allocated
    }

    /// Check if garbage collection should be triggered
    pub fn should_collect(&self) -> bool {
        self.bytes_allocated >= self.next_gc
    }

    /// Perform mark-and-sweep garbage collection
    pub fn collect(&mut self, roots: &[Value]) {
        self.stats.collections += 1;

        // Phase 1: Mark - trace all reachable objects
        let mut tracer = Tracer::new();

        // Mark all root objects
        for root in roots {
            mark_value(root, &mut tracer, &self.objects);
        }

        // Process grey set until empty
        while let Some(id) = tracer.pop_grey() {
            if let Some(obj) = self.objects.get(&id) {
                mark_value(&obj.value, &mut tracer, &self.objects);
            }
        }

        // Phase 2: Sweep - collect unmarked objects
        let mut to_remove = Vec::new();
        let mut bytes_freed = 0;

        for (id, obj) in &self.objects {
            if !tracer.marked.contains(id) {
                to_remove.push(*id);
                bytes_freed += estimate_value_size(&obj.value);
            }
        }

        // Remove unmarked objects
        for id in to_remove {
            self.objects.remove(&id);
            self.stats.objects_collected += 1;
        }

        self.bytes_allocated = self.bytes_allocated.saturating_sub(bytes_freed);
        self.stats.bytes_collected += bytes_freed;

        // Reset all remaining objects to white for next cycle
        for obj in self.objects.values_mut() {
            obj.header.reset();
        }

        // Adjust GC threshold based on live set size
        self.next_gc = (self.bytes_allocated * 2).max(1024 * 1024);
    }
}

impl Default for GcHeap {
    fn default() -> Self {
        Self::new()
    }
}

/// Mark a value and all values it references
fn mark_value(value: &Value, tracer: &mut Tracer, objects: &HashMap<usize, GcObject>) {
    match value {
        Value::Cons { head, tail } => {
            mark_value(head, tracer, objects);
            mark_value(tail, tracer, objects);
        }
        Value::Tuple(elements) => {
            for elem in elements {
                mark_value(elem, tracer, objects);
            }
        }
        Value::Array(arr) => {
            let arr = arr.borrow();
            for elem in arr.iter() {
                mark_value(elem, tracer, objects);
            }
        }
        Value::Record(fields) => {
            let fields = fields.borrow();
            for value in fields.values() {
                mark_value(value, tracer, objects);
            }
        }
        Value::Map(map) => {
            let map = map.borrow();
            for value in map.values() {
                mark_value(value, tracer, objects);
            }
        }
        Value::Variant { fields, .. } => {
            for field in fields {
                mark_value(field, tracer, objects);
            }
        }
        Value::Closure(closure) => {
            mark_closure(closure, tracer, objects);
        }
        Value::NativeFn { args, .. } => {
            for arg in args {
                mark_value(arg, tracer, objects);
            }
        }
        // Primitive types don't need tracing
        Value::Int(_) | Value::Float(_) | Value::Bool(_) | Value::Str(_) | Value::Unit | Value::Nil => {}
        // HostData is managed by Rust's reference counting
        Value::HostData(_) => {}
    }
}

/// Mark a closure and all values it captures
fn mark_closure(closure: &Closure, tracer: &mut Tracer, objects: &HashMap<usize, GcObject>) {
    // Mark constants in the closure's chunk
    for constant in &closure.chunk.constants {
        mark_value(constant, tracer, objects);
    }

    // Mark upvalues
    for upvalue in &closure.upvalues {
        let upvalue = upvalue.borrow();
        match &*upvalue {
            Upvalue::Closed(value) => mark_value(value, tracer, objects),
            Upvalue::Open(_) => {
                // Open upvalues reference stack slots, which are roots themselves
            }
        }
    }
}

/// Estimate the size of a value in bytes (approximate)
fn estimate_value_size(value: &Value) -> usize {
    match value {
        Value::Int(_) => 8,
        Value::Float(_) => 8,
        Value::Bool(_) => 1,
        Value::Str(s) => std::mem::size_of::<String>() + s.len(),
        Value::Unit => 0,
        Value::Tuple(elements) => {
            std::mem::size_of::<Vec<Value>>() + elements.iter().map(estimate_value_size).sum::<usize>()
        }
        Value::Cons { head, tail } => {
            16 + estimate_value_size(head) + estimate_value_size(tail)
        }
        Value::Nil => 0,
        Value::Array(arr) => {
            let arr = arr.borrow();
            std::mem::size_of::<Vec<Value>>() + arr.iter().map(estimate_value_size).sum::<usize>()
        }
        Value::Record(fields) => {
            let fields = fields.borrow();
            std::mem::size_of::<HashMap<String, Value>>()
                + fields
                    .iter()
                    .map(|(k, v)| k.len() + estimate_value_size(v))
                    .sum::<usize>()
        }
        Value::Map(map) => {
            let map = map.borrow();
            std::mem::size_of::<HashMap<String, Value>>()
                + map
                    .iter()
                    .map(|(k, v)| k.len() + estimate_value_size(v))
                    .sum::<usize>()
        }
        Value::Variant { type_name, variant_name, fields } => {
            type_name.len()
                + variant_name.len()
                + fields.iter().map(estimate_value_size).sum::<usize>()
        }
        Value::Closure(c) => {
            // Rough estimate: chunk size + upvalues
            std::mem::size_of::<Closure>()
                + c.chunk.instructions.len() * 4
                + c.upvalues.len() * 8
        }
        Value::NativeFn { name, args, .. } => {
            name.len() + args.iter().map(estimate_value_size).sum::<usize>()
        }
        Value::HostData(_) => 8, // Just the Rc pointer
    }
}

// Implement Trace for standard containers
impl<T: Trace> Trace for Vec<T> {
    fn trace(&self, tracer: &mut Tracer) {
        for item in self {
            item.trace(tracer);
        }
    }
}

impl<K, V: Trace> Trace for HashMap<K, V> {
    fn trace(&self, tracer: &mut Tracer) {
        for value in self.values() {
            value.trace(tracer);
        }
    }
}

impl Trace for Value {
    fn trace(&self, tracer: &mut Tracer) {
        match self {
            Value::Cons { head, tail } => {
                head.trace(tracer);
                tail.trace(tracer);
            }
            Value::Tuple(elements) => elements.trace(tracer),
            Value::Array(arr) => {
                let arr = arr.borrow();
                for elem in arr.iter() {
                    elem.trace(tracer);
                }
            }
            Value::Record(fields) => {
                let fields = fields.borrow();
                for value in fields.values() {
                    value.trace(tracer);
                }
            }
            Value::Variant { fields, .. } => fields.trace(tracer),
            Value::Closure(closure) => closure.trace(tracer),
            Value::NativeFn { args, .. } => args.trace(tracer),
            _ => {}
        }
    }
}

impl Trace for Closure {
    fn trace(&self, tracer: &mut Tracer) {
        // Trace constants in chunk
        for constant in &self.chunk.constants {
            constant.trace(tracer);
        }

        // Trace upvalues
        for upvalue in &self.upvalues {
            let upvalue = upvalue.borrow();
            match &*upvalue {
                Upvalue::Closed(value) => value.trace(tracer),
                Upvalue::Open(_) => {
                    // Open upvalues reference stack slots
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::rc::Rc;
    use std::cell::RefCell;
    use super::*;

    #[test]
    fn test_gc_header_creation() {
        let header = GcHeader::new(1);
        assert_eq!(header.id, 1);
        assert_eq!(header.color, Color::White);
        assert!(!header.is_marked());
    }

    #[test]
    fn test_gc_header_marking() {
        let mut header = GcHeader::new(1);
        header.mark();
        assert_eq!(header.color, Color::Grey);
        assert!(header.is_marked());
    }

    #[test]
    fn test_gc_header_reset() {
        let mut header = GcHeader::new(1);
        header.mark();
        header.reset();
        assert_eq!(header.color, Color::White);
        assert!(!header.is_marked());
    }

    #[test]
    fn test_tracer_mark() {
        let mut tracer = Tracer::new();
        tracer.mark(1);
        assert!(tracer.marked.contains(&1));
        assert_eq!(tracer.grey_set.len(), 1);
    }

    #[test]
    fn test_tracer_mark_duplicate() {
        let mut tracer = Tracer::new();
        tracer.mark(1);
        tracer.mark(1);
        assert_eq!(tracer.marked.len(), 1);
        assert_eq!(tracer.grey_set.len(), 1);
    }

    #[test]
    fn test_tracer_pop_grey() {
        let mut tracer = Tracer::new();
        tracer.mark(1);
        tracer.mark(2);

        assert!(tracer.has_grey());
        let id = tracer.pop_grey();
        assert_eq!(id, Some(2));

        let id = tracer.pop_grey();
        assert_eq!(id, Some(1));

        assert!(!tracer.has_grey());
    }

    #[test]
    fn test_gc_heap_allocate() {
        let mut heap = GcHeap::new();
        let id = heap.allocate(Value::Int(42));

        assert_eq!(heap.object_count(), 1);
        assert_eq!(heap.get(id), Some(&Value::Int(42)));
    }

    #[test]
    fn test_gc_heap_multiple_allocations() {
        let mut heap = GcHeap::new();
        let id1 = heap.allocate(Value::Int(1));
        let id2 = heap.allocate(Value::Int(2));
        let id3 = heap.allocate(Value::Int(3));

        assert_eq!(heap.object_count(), 3);
        assert_eq!(heap.get(id1), Some(&Value::Int(1)));
        assert_eq!(heap.get(id2), Some(&Value::Int(2)));
        assert_eq!(heap.get(id3), Some(&Value::Int(3)));
    }

    #[test]
    fn test_gc_collect_no_roots() {
        let mut heap = GcHeap::new();
        heap.allocate(Value::Int(1));
        heap.allocate(Value::Int(2));
        heap.allocate(Value::Int(3));

        assert_eq!(heap.object_count(), 3);

        // Collect with no roots - all should be collected
        heap.collect(&[]);

        assert_eq!(heap.object_count(), 0);
        assert_eq!(heap.stats.objects_collected, 3);
    }

    #[test]
    fn test_gc_collect_with_roots() {
        let mut heap = GcHeap::new();
        let id1 = heap.allocate(Value::Int(1));
        let _id2 = heap.allocate(Value::Int(2));
        let id3 = heap.allocate(Value::Int(3));

        // Keep id1 and id3 as roots
        let roots = vec![Value::Int(1), Value::Int(3)];

        heap.collect(&roots);

        // Since our simple implementation doesn't actually track object IDs in values,
        // this test demonstrates the collection mechanism
        assert_eq!(heap.stats.collections, 1);
    }

    #[test]
    fn test_gc_collect_list() {
        let mut heap = GcHeap::new();

        // Create a list [1, 2, 3]
        let list = Value::vec_to_cons(vec![Value::Int(1), Value::Int(2), Value::Int(3)]);

        // Allocate it (in a real implementation, the list nodes would be on the heap)
        heap.allocate(list.clone());

        // With the list as root, it should survive collection
        heap.collect(&[list]);

        assert_eq!(heap.stats.collections, 1);
    }

    #[test]
    fn test_gc_collect_cycle_simple() {
        let mut heap = GcHeap::new();

        // Create two records that reference each other (cycle)
        let mut fields1 = HashMap::new();
        fields1.insert("value".to_string(), Value::Int(1));
        let record1 = Value::Record(Rc::new(RefCell::new(fields1)));

        let mut fields2 = HashMap::new();
        fields2.insert("value".to_string(), Value::Int(2));
        fields2.insert("next".to_string(), record1.clone());
        let record2 = Value::Record(Rc::new(RefCell::new(fields2)));

        // Create the cycle
        if let Value::Record(r1) = &record1 {
            r1.borrow_mut().insert("next".to_string(), record2.clone());
        }

        heap.allocate(record1.clone());
        heap.allocate(record2.clone());

        // Without roots, the cycle should be collected
        heap.collect(&[]);

        assert_eq!(heap.stats.objects_collected, 2);
    }

    #[test]
    fn test_estimate_value_size() {
        assert_eq!(estimate_value_size(&Value::Int(42)), 8);
        assert_eq!(estimate_value_size(&Value::Bool(true)), 1);
        assert_eq!(estimate_value_size(&Value::Unit), 0);
        assert_eq!(estimate_value_size(&Value::Nil), 0);

        let str_val = Value::Str("hello".to_string());
        assert!(estimate_value_size(&str_val) > 5);

        let tuple = Value::Tuple(vec![Value::Int(1), Value::Int(2)]);
        assert!(estimate_value_size(&tuple) > 16);
    }

    #[test]
    fn test_gc_threshold() {
        let threshold = 1024;
        let mut heap = GcHeap::with_threshold(threshold);

        // Allocate until we exceed threshold
        while !heap.should_collect() {
            heap.allocate(Value::Int(42));
        }

        assert!(heap.bytes_allocated() >= threshold);
    }

    #[test]
    fn test_gc_stats() {
        let mut heap = GcHeap::new();
        heap.allocate(Value::Int(1));
        heap.allocate(Value::Int(2));

        assert_eq!(heap.stats.collections, 0);

        heap.collect(&[]);

        assert_eq!(heap.stats.collections, 1);
        assert!(heap.stats.objects_collected > 0);
    }
}
