from functools import partial

import numpy as np
import pytest
from PIL import Image
from funcy import lmap, compose, flatten

from pyinfra.rest import pack, normalize_item
from pyinfra.utils.func import star, lift
from test.utils.image import image_to_bytes


@pytest.fixture
def data(data_type, pdf):
    if data_type == "pdf":
        return pdf
    elif data_type == "bytestring":
        return "content".encode("latin1")


@pytest.fixture
def input_data_items(item_type, n_items, pdf):
    if item_type == "string":
        return [bytes(f"content{i}", encoding="utf8") for i in range(n_items)]
    elif item_type == "image":
        return images(n_items)
    elif item_type == "pdf":
        return [pdf] * n_items
    else:
        raise ValueError(f"Unknown item type {item_type}")


@pytest.fixture
def target_data_items(input_data_items, item_type, operation, metadata):
    op = compose(lift(star(pack)), normalize_item, partial(operation, metadata=metadata))
    expected = list(flatten(map(op, input_data_items)))
    return expected


@pytest.fixture(params=[0, 1, 2, 5])
def n_items(request):
    return request.param


def array_to_image(array) -> Image.Image:
    return Image.fromarray(np.uint8(array * 255), mode="RGB")


def input_batch(n_items):
    return np.random.random_sample(size=(n_items, 3, 30, 30))


def images(n_items):
    return lmap(compose(image_to_bytes, array_to_image), input_batch(n_items))


@pytest.fixture
def metadata():
    return {"idx": [1, 100, 101], "path": "asd/asd"}