import io
import socket
from multiprocessing import Process

import fitz
import pytest
import requests
from PIL import Image
from funcy import retry
from waitress import serve

from pyinfra.server.server import set_up_processing_server, make_streamable_and_wrap_in_packing_logic
from pyinfra.utils.func import starlift
from test.utils.image import image_to_bytes


@pytest.fixture
def host():
    return "0.0.0.0"


def get_free_port(host):
    sock = socket.socket()
    sock.bind((host, 0))
    return sock.getsockname()[1]


@pytest.fixture
def port(host):
    return get_free_port(host)


@pytest.fixture
def url(host, port):
    return f"http://{host}:{port}"


@pytest.fixture
def server(processor_fn):
    return set_up_processing_server(processor_fn)


@pytest.fixture
def processor_fn(operation_conditionally_batched, buffer_size, batched):
    return make_streamable_and_wrap_in_packing_logic(operation_conditionally_batched, batched, buffer_size)


@pytest.fixture
def operation_conditionally_batched(operation, batched):
    return starlift(operation) if batched else operation


@pytest.fixture
def operation(item_type, batched, one_to_many):
    def upper(string: bytes, metadata):
        return string.decode().upper().encode(), metadata

    def duplicate(string: bytes, metadata):
        for _ in range(2):
            yield upper(string, metadata)

    def rotate(im: bytes, metadata):
        im = Image.open(io.BytesIO(im))
        return image_to_bytes(im.rotate(90)), metadata

    def stream_pages(pdf: bytes, metadata):
        for i, page in enumerate(fitz.open(stream=pdf)):
            # yield page.get_pixmap().tobytes("png"), metadata
            yield f"page_{i}".encode(), metadata

    try:
        return {"string": duplicate if one_to_many else upper, "image": rotate, "pdf": stream_pages}[item_type]
    except KeyError:
        raise ValueError(f"No operation specified for item type {item_type}")


@pytest.fixture(params=["pdf", "string", "image"])
def item_type(request):
    return request.param


@pytest.fixture(params=[True, False])
def one_to_many(request):
    return request.param


@pytest.fixture(params=[False, True])
def batched(request):
    """Controls, whether the buffer processor function of the webserver is applied to batches or single items."""
    return request.param


@pytest.fixture
def host_and_port(host, port):
    return {"host": host, "port": port}


@retry(tries=5, timeout=1)
def server_ready(url):
    response = requests.get(f"{url}/ready")
    response.raise_for_status()
    return response.status_code == 200


@pytest.fixture(autouse=False, scope="function")
def server_process(server, host_and_port, url):
    def get_server_process():
        return Process(target=serve, kwargs={"app": server, **host_and_port})

    server = get_server_process()
    server.start()

    if server_ready(url):
        yield

    server.kill()
    server.join()
    server.close()