from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw


SUPPORTED_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}


class ImageState:
    def __init__(self, assets_dir: str = "assets"):
        self.assets_dir = Path(assets_dir)
        self.image_paths = self._load_image_paths()
        self.index = 0

    def _load_image_paths(self):
        if not self.assets_dir.exists():
            return []

        paths = []
        for path in sorted(self.assets_dir.iterdir()):
            if path.suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS:
                paths.append(path)
        return paths

    def next_image(self):
        if not self.image_paths:
            return None
        self.index = (self.index + 1) % len(self.image_paths)
        return self.image_paths[self.index]

    def prev_image(self):
        if not self.image_paths:
            return None
        self.index = (self.index - 1) % len(self.image_paths)
        return self.image_paths[self.index]

    def first_image(self):
        if not self.image_paths:
            return None
        self.index = 0
        return self.image_paths[self.index]

    def last_image(self):
        if not self.image_paths:
            return None
        self.index = len(self.image_paths) - 1
        return self.image_paths[self.index]

    def get_current_image(self):
        if not self.image_paths:
            return self._create_placeholder_image()
        image = Image.open(self.image_paths[self.index]).convert("RGB")
        return np.array(image)

    def _create_placeholder_image(self):
        image = Image.new("RGB", (640, 360), color=(40, 40, 40))
        draw = ImageDraw.Draw(image)
        draw.text((20, 20), "No images in assets/", fill=(255, 255, 255))
        return np.array(image)