import io
import socket
from multiprocessing import Process
from operator import attrgetter

import fitz
import pytest
import requests
from PIL import Image
from funcy import retry, compose, flatten
from waitress import serve

from pyinfra.rest import unpack_op_pack, unpack_batchop_pack, inspect
from pyinfra.utils.buffer import bufferize
from pyinfra.utils.func import llift, starlift
from test.server import set_up_processing_server
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, buffer_size, batched):

    if batched:
        operation = starlift(operation)

    wrapper = unpack_batchop_pack if batched else compose(llift, unpack_op_pack)
    operation = wrapper(operation)

    return bufferize(operation, buffer_size=buffer_size, persist_fn=attrgetter("json"))


@pytest.fixture
def operation(item_type, batched):
    def upper(string: bytes, metadata):
        return string.decode().upper().encode(), 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 page in fitz.open(stream=pdf):
            yield page.get_pixmap().tobytes("png"), metadata

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


@pytest.fixture(params=["string"])
def item_type(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(params=[1, 3, 7, 100])
def buffer_size(request):
    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()