from interface import PCIeInterface
import threading
import queue
import time

class NVMeCommand:
    def __init__(self, cmd_type, lba, data=None):
        self.cmd_type = cmd_type  # 'read' or 'write'
        self.lba = lba
        self.data = data
        self.result = None
        self.completed = threading.Event()

class NVMeController:
    def __init__(self, ssd_controller, queue_depth=64):
        self.ssd = ssd_controller
        self.submission_queue = queue.Queue(maxsize=queue_depth)
        self.completion_queue = queue.Queue(maxsize=queue_depth)
        self.running = True
        self.worker = threading.Thread(target=self.process_commands)
        self.worker.daemon = True
        self.worker.start()
        self.interface = PCIeInterface()

    def submit(self, cmd):
        self.submission_queue.put(cmd)

    def process_commands(self):
        while self.running:
            try:
                cmd = self.submission_queue.get(timeout=0.1)
                if cmd.cmd_type == 'write':
                    self.ssd.program(cmd.lba, cmd.data)
                    self.interface.simulate_transfer(len(cmd.data) * 32 // 8, direction='write')
                    cmd.result = 'write_complete'
                elif cmd.cmd_type == 'read':
                    data = self.ssd.read(cmd.lba)
                    self.interface.simulate_transfer(len(data) * 32 // 8, direction='read')
                    cmd.result = data
                self.completion_queue.put(cmd)
                cmd.completed.set()
            except queue.Empty:
                continue

    def get_completion(self, timeout=1.0):
        try:
            cmd = self.completion_queue.get(timeout=timeout)
            return cmd
        except queue.Empty:
            return None

    def shutdown(self):
        self.running = False
        self.worker.join()