#!/usr/bin/env python3
# examples/basic_operations.py
# 
# This example demonstrates basic NAND flash operations including:
# - Initialization
# - Reading
# - Writing
# - Erasing
# - Error handling
# - Shutdown

import os
import sys
import time

# Add the project root directory to the Python path
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir)
sys.path.insert(0, project_root)

from src.opennandlab.simulator import NANDController
from src.opennandlab.config import SimulatorConfig


def basic_operations_example():
    """
    Demonstrate basic operations with NAND flash controller
    """
    print("=== Basic NAND Flash Operations Example ===")
    
    # Load default configuration
    config = SimulatorConfig()
    
    # Create NAND controller (simulation mode for safety)
    controller = NANDController(config, simulation_mode=True)
    print("NAND controller created")
    
    try:
        # Initialize the controller
        print("\n--- Initializing NAND Controller ---")
        controller.initialize()
        print("NAND controller initialized successfully")
        
        # Get and show device information
        device_info = controller.get_device_info()
        print("\n--- NAND Device Information ---")
        print(f"Page Size: {device_info['config']['page_size']} bytes")
        print(f"Block Size: {device_info['config']['block_size']} pages")
        print(f"Number of Blocks: {device_info['config']['num_blocks']}")
        
        # Writing and reading use logical block numbers (lbn) in v2.0.0
        # The FTL manages the physical translation.
        lbn = 10  # A random logical block number
        
        # Write to the logical page
        print("\n--- Writing Data ---")
        test_data = f"Test data written to logical page {lbn} at {time.time()}".encode('utf-8')
        controller.write_page(lbn, test_data)
        print(f"Data written to logical page {lbn}")
        
        # Read from the logical page
        print("\n--- Reading Data ---")
        read_data = controller.read_page(lbn)
        print(f"Read {len(read_data)} bytes from logical page {lbn}")
        
        # Verify the data
        if test_data in read_data:
            print("Data verification successful!")
            print(f"Original: {test_data}")
            print(f"Read: {read_data[:len(test_data)]}")
        else:
            print("Data verification failed!")
            print(f"Original: {test_data}")
            print(f"Read: {read_data[:100]}")
            
        # Demonstrate error handling
        print("\n--- Error Handling Example ---")
        try:
            # Try to access an invalid logical page (beyond range)
            invalid_lbn = controller.ftl.num_logical_pages + 10
            controller.read_page(invalid_lbn)
        except Exception as e:
            print(f"Expected error caught: {e}")
        
        # Get usage statistics
        print("\n--- NAND Usage Statistics ---")
        device_info = controller.get_device_info()
        stats = device_info.get('statistics', {})
        
        print(f"Total reads: {stats.get('reads', 0)}")
        print(f"Total writes: {stats.get('writes', 0)}")
        print(f"Total erases: {stats.get('erases', 0)}")
        
        if 'bad_blocks' in stats:
            bad_block_stats = stats['bad_blocks']
            print(f"Bad blocks: {bad_block_stats.get('count', 0)} ({bad_block_stats.get('percentage', 0):.2f}%)")
        
        if 'wear_leveling' in stats:
            wear_stats = stats['wear_leveling']
            print(f"Min erase count: {wear_stats.get('min_erase_count', 0)}")
            print(f"Max erase count: {wear_stats.get('max_erase_count', 0)}")
            print(f"Avg erase count: {wear_stats.get('avg_erase_count', 0):.2f}")
            
    except Exception as e:
        print(f"Error during operations: {e}")
    finally:
        # Always shut down the controller properly
        print("\n--- Shutting Down NAND Controller ---")
        controller.shutdown()
        print("NAND controller shut down successfully")

if __name__ == "__main__":
    basic_operations_example()