import math
import os
from enum import Enum
from io import BytesIO
from typing import Any, Dict

import faker
from faker_file.providers.png_file import GraphicPngFileProvider
from faker_file.storages.filesystem import FileSystemStorage
from PIL import Image

from deltacat.benchmarking.data.row_generator import RowGenerator


class ImageStyle(Enum):
    RANDOM_BYTES = 1
    """Generate random bytes to simulate an image.

    This is the fastest option (if you want to test correctness).
    """
    PILLOW = 2
    """Generate actual PNG files in-memory directly using Pillow"""
    FAKER_FILE = 3
    """Generate PNG files on-disk using with some random elements"""


class RandomRowGenerator(RowGenerator):
    """Generate rows with 'images' that are just randomly-generated bytes"""

    def __init__(
        self, seed=0, tmp_dir=None, style: ImageStyle = ImageStyle.RANDOM_BYTES
    ):
        self.seed = seed
        self.fake = faker.Faker()
        self.fake.seed_instance(seed)
        self.fake.add_provider(GraphicPngFileProvider)
        self.temp_dir = tmp_dir
        self.style = style

    def __str__(self):
        return f"random source"

    def _generate_image(self, width, height) -> bytes:
        if self.style == ImageStyle.RANDOM_BYTES:
            return self._generate_with_random_bytes(width, height)
        elif self.style == ImageStyle.PILLOW:
            return self._generate_with_pillow(width, height)
        elif self.style == ImageStyle.FAKER_FILE:
            return self._generate_with_faker(width, height)
        else:
            raise ValueError("Unknown ImageStyle")

    @staticmethod
    def _generate_with_random_bytes(width, height) -> bytes:
        """Generate random bytes to simulate an image."""
        target_size = math.floor(
            width * height / 50
        )  # this isn't actually how file size relates to image size
        # Assumption: we don't actually need images. It suffices to generate arbitrary-length bytes of random characters.
        return os.urandom(target_size)

    @staticmethod
    def _generate_with_pillow(width, height) -> bytes:
        """Generate actual PNG files in-memory directly using Pillow"""
        file = BytesIO()
        image = Image.new("RGBA", size=(width, height), color=(155, 0, 0))
        image.save(file, "png")
        file.name = "test.png"
        file.seek(0)
        return file.read()

    def _generate_with_faker(self, width, height) -> bytes:
        """Generate PNG files on-disk using with some random elements"""
        rel_name = self.fake.graphic_png_file(
            storage=FileSystemStorage(
                root_path=self.temp_dir,
                rel_path="tmp",
            ),
            size=(width, height),
        )
        file_name = f"{self.temp_dir}/{rel_name}"
        with open(file_name, "rb") as f:
            return f.read()

    def generate(self) -> Dict[str, Any]:
        return {
            "id": self.fake.random_int(0, 10_000_000),
            "source": self.fake.image_url(),
            "media": (
                self._generate_image(
                    self.fake.random_int(512, 2048), self.fake.random_int(512, 4096)
                )
            ),
        }
