import utime


class Rainbow:
    def __init__(self, driver):
        self.driver = driver

    def _wheel(self, pos):
        if pos < 85:
            return (pos * 3, 255 - pos * 3, 0)
        elif pos < 170:
            pos -= 85
            return (255 - pos * 3, 0, pos * 3)
        else:
            pos -= 170
            return (0, pos * 3, 255 - pos * 3)

    def run(self, preset):
        step = self.driver.step % 256
        step_amount = max(1, int(preset.n1))  # n1 controls step increment

        # If auto is False, run a single step and then stop
        if not preset.a:
            for i in range(self.driver.num_leds):
                rc_index = (i * 256 // self.driver.num_leds) + step
                self.driver.n[i] = self.driver.apply_brightness(self._wheel(rc_index & 255), preset.b)
            self.driver.n.write()
            # Increment step by n1 for next manual call
            self.driver.step = (step + step_amount) % 256
            # Allow tick() to advance the generator once
            yield
            return

        last_update = utime.ticks_ms()

        while True:
            current_time = utime.ticks_ms()
            sleep_ms = max(1, int(preset.d))  # Get delay from preset
            if utime.ticks_diff(current_time, last_update) >= sleep_ms:
                for i in range(self.driver.num_leds):
                    rc_index = (i * 256 // self.driver.num_leds) + step
                    self.driver.n[i] = self.driver.apply_brightness(
                        self._wheel(rc_index & 255),
                        preset.b,
                    )
                self.driver.n.write()
                step = (step + step_amount) % 256
                self.driver.step = step
                last_update = current_time
            # Yield once per tick so other logic can run
            yield