import micropython
import gc
from pybricks import version
class MicroPythonDiagnostics:
    def __init__(self, hub):
        self.successfultests = 0
        self.failedtests = {}
    def testgcmanual(self):
        gc.disable()
        print(f"Initial free: {gc.mem_free()} bytes")
        large_data = [i for i in range(10000)] 
        print(f"After allocation: {gc.mem_free()} bytes")
        gc.collect()
        aftergcstillref = gc.mem_free()
        print(f"After gc.collect (data still referenced): {aftergcstillref} bytes")
        large_data = None 
        print("Reference to data removed.")
        aftergcnoref = gc.mem_free()
        gc.collect()
        print(f"After gc.collect (data dead): {aftergcnoref} bytes")
        if(aftergcnoref < aftergcstillref):
            print("Completed Test 4/5: Manual garbage collection - SUCCESSFUL")
            self.successfultests += 1
        else:
            print("Completed Test 4/5: Manual garbage collection - FAILED")
            self.failedtests["Manual garbage collection"] = "Heap not cleared"
    def testgcauto(self):
        input("Disabling garbage collection. The amount of used memory should quickly increase. Press Enter to begin:")
        gc.disable()
        gc.threshold(5000) 
        
        total_mem = 255616
        bytes_per_hash = 3000
        
        print("Memory Monitor: [# = Used] [. = Free]")
        print("-" * (total_mem // bytes_per_hash))

        for i in range(500):
            _ = bytearray(300) 
            
            if i % 25 == 0:
                used = gc.mem_alloc()
                free = gc.mem_free()

                hashes = used // bytes_per_hash
                dots = free // bytes_per_hash
                
                print(f"{i:03d}: [{'#' * hashes}{'.' * dots}] {free} bytes free")
        final_disabled_free = gc.mem_free()
        input("Enabling garbage collection. The amount of used memory should stay relatively low. Press Enter to begin:")
        gc.enable()
        gc.threshold(5000) 
        
        total_mem = 255616
        bytes_per_hash = 3000
        
        print("Memory Monitor: [# = Used] [. = Free]")
        print("-" * (total_mem // bytes_per_hash))

        for i in range(500):
            _ = bytearray(300) 

            if i % 25 == 0:
                used = gc.mem_alloc()
                free = gc.mem_free()
                
                hashes = used // bytes_per_hash
                dots = free // bytes_per_hash

                print(f"{i:03d}: [{'#' * hashes}{'.' * dots}] {free} bytes free")

        final_enabled_free = gc.mem_free()
        if final_enabled_free > final_disabled_free:
            print("Completed Test 5/5: Automatic garbage collection - SUCCESSFUL")
            print(f"Difference: {final_enabled_free - final_disabled_free} bytes saved.")
            self.successfultests += 1
        else:
            print("Completed Test 5/5: Automatic garbage collection - FAILED")
            self.failedtests["Automatic garbage collection"] = "No GC difference"
    def performMemoryDiagnostics(self):
        input("Press Enter to retrieve memory information:")
        print("[Hub Diagnostics - MicroPython - Memory] Memory information (retrieved from the MicroPython environment):")
        micropython.mem_info(1)
        input("After you're done reading the results, press Enter to run heap diagnostics:")
        print("[Hub Diagnostics - MicroPython - Memory] Testing heap lock and unlock.")
        print("[Hub Diagnostics - MicroPython - Memory] Allocating memory while heap is unlocked:")
        try:
            x = [1, 2, 3, 4, 5]
            print("[Hub Diagnostics - MicroPython - Memory] Completed Test 1/5: Normal memory allocation - SUCCESS")
            print("There was no MemoryError raised. The value of the new variable x is", x)
            self.successfultests += 1
        except MemoryError:
            print("[Hub Diagnostics - MicroPython - Memory] Completed Test 1/5: Normal memory allocation - FAILED")
            self.failedtests["Normal memory allocation"] = "MemoryError"
        print("[Hub Diagnostics - MicroPython - Memory] Locking the heap:")
        micropython.heap_lock()
        print("[Hub Diagnostics - MicroPython - Memory] Heap was locked. Attempting to allocate memory (this should fail):")
        try:
            y = [10, 20, 30, 40, 50]
            print("[Hub Diagnostics - MicroPython - Memory] Completed Test 2/5: Heap lock - FAILED")
            self.failedtests["Heap lock"] = "No heap lock"
        except MemoryError:
            print("[Hub Diagnostics - MicroPython - Memory] Completed Test 2/5: Heap lock - SUCCESS")
            self.successfultests += 1
        print("[Hub Diagnostics - MicroPython - Memory] Unlocking the heap:")
        micropython.heap_unlock()
        print("[Hub Diagnostics - MicroPython - Memory] Heap was unlocked. Attempting to allocate memory (this should succeed):")
        try:
            z = [100, 200, 300, 400, 500]
            print("[Hub Diagnostics - MicroPython - Memory] Completed Test 3/5: Heap unlock - FAILED")
            print("The value of the new variable z is", z)
            self.successfultests += 1
        except MemoryError:
            print("[Hub Diagnostics - MicroPython - Memory] Completed Test 3/5: Heap unlock - FAILED")
            self.failedtests["Heap unlock"] = "No heap unlock"
    def printAll(self):
        self.performMemoryDiagnostics()
        input("After you're done reading the results, press Enter to run manual garbage collection test:")
        self.testgcmanual()
        input("After you're done reading the results, press Enter to run automatic garbage collection test:")
        self.testgcauto()
        print(f"\n=== Results: {self.successfultests}/5 tests passed ===")
        if self.failedtests:
            print("Failed tests:")
            for key, value in self.failedtests.items():
                print(f"  {key}: {value}")
#test = MicroPythonDiagnostics(hub=PrimeHub())
#test.printAll()