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
from waitress import serve

from pyinfra.server.utils import unpack_op_pack, unpack_batchop_pack
from pyinfra.utils.buffer import bufferize
from pyinfra.utils.func import llift, starlift
from pyinfra.server.server import set_up_processing_server, make_processor_fn
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):
    return make_processor_fn(operation, buffer_size, batched)


@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 i, page in enumerate(fitz.open(stream=pdf)):
            # yield page.get_pixmap().tobytes("png"), metadata
            yield f"page_{i}".encode(), 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
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()