import io
import sys
import tarfile
import docker
import pytest
import requests
import subprocess
import time
import uuid

from contextlib import contextmanager
from pathlib import Path
from urllib.parse import quote
from loguru import logger


IMAGES_ENDPOINT = 'http://localhost:5000/api/v1.0/images'
BROKER_ENDPOINT = 'http://guest:guest@localhost:15672/api'
TASK_QUEUE_ENDPOINT = f'{BROKER_ENDPOINT}/queues/%2F/task_queue'

logger.remove()
logger.add(sys.stderr, colorize=False, format="=== TEST ===| {time:YYYY-MM-DD HH:mm:ss.SSS} {level} {message}")


# Tests ===============================================================================================================

@pytest.mark.parametrize("services", [['rabbitmq', 'server']])
def test_empty_data_dir(docker_tester):
    logger.info(f"Sending GET {IMAGES_ENDPOINT}")
    try:
        response = requests.get(IMAGES_ENDPOINT, timeout=3)
    except Exception as e:
        logger.error(f"Request failed: {e}")
        pytest.fail(f"Failed to get images: {e}")
    logger.info(f"Got response: {response.status_code} {response.text.rstrip()}")
    assert response.status_code == 200
    assert 'image_ids' in response.json()
    assert len(response.json()['image_ids']) == 0


@pytest.mark.parametrize("services", [['rabbitmq', 'server']])
def test_bad_request(docker_tester):
    invalid_bodies = [{}, {'image_url': None}, {'image_url': 1},
                      {'image_url': ''}, [], None, 'image-url']
    for body in invalid_bodies:
        logger.info(f"Sending invalid POST: {body!r}")
        response = requests.post(IMAGES_ENDPOINT, json=body, timeout=3)
        assert response.status_code == 400, response.text
    response = requests.post(IMAGES_ENDPOINT, data='{',
                             headers={'Content-Type': 'application/json'}, timeout=3)
    assert response.status_code == 400, response.text


@pytest.mark.parametrize("services", [['rabbitmq', 'server']])
def test_nonexistent_image(docker_tester):
    nonexistent_image_id = str(uuid.uuid4())
    logger.info(f"Sending GET {IMAGES_ENDPOINT}/{nonexistent_image_id}")
    try:
        response = requests.get(f'{IMAGES_ENDPOINT}/{nonexistent_image_id}', timeout=3)
    except Exception as e:
        logger.error(f"Request failed: {e}")
        pytest.fail(f"Failed to get images: {e}")
    logger.info(f"Got response: {response.status_code} {response.text.rstrip()}")
    assert response.status_code == 404


@pytest.mark.parametrize("services", [['rabbitmq', 'server']])
def test_task_queue(docker_tester):
    time.sleep(5)
    check_task_queue(0, 10)
    pending_ids = post_images(10)
    check_unprocessed_images(pending_ids, docker_tester)
    check_task_queue(10, 10)
    time.sleep(5)
    check_task_queue(10, 10)
    check_unprocessed_images(pending_ids, docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_single_image(docker_tester):
    pending_ids = post_images(1)
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_multiple_images(docker_tester):
    pending_ids = post_images(10)
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_captions_generated_on_workers(docker_tester):
    worker1 = docker_tester.containers.get("distsys-mq-worker-1")
    worker2 = docker_tester.containers.get("distsys-mq-worker-2")
    wait_for_worker_ready(worker1)
    wait_for_worker_ready(worker2)
    worker1.pause()
    worker2.pause()

    pending_ids = post_images(10)
    time.sleep(5)

    check_unprocessed_images(pending_ids, docker_tester)
    worker1.unpause()
    worker2.unpause()
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server-fdv', 'worker']])
def test_multiple_images_no_listdir(docker_tester):
    pending_ids = post_images(10)
    # This server has no access to the workers' result volume.
    wait_and_check_results(pending_ids, 10, check_captions=False, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_heartbeats_timeout(docker_tester):
    # With heartbeat=5, closing a connection without heartbeats can take 15s.
    time.sleep(20)
    pending_ids = post_images(10)
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_publisher_confirms(docker_tester):
    rabbit = docker_tester.containers.get("distsys-mq-rabbitmq-1")
    rabbit.pause()
    pending_ids = post_images(10)
    time.sleep(5)
    rabbit.kill()
    time.sleep(1)
    rabbit.start()
    wait_for_task_queue(lambda queue: True, "broker restart", max_attempts=20)
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_faulty_worker(docker_tester):
    worker1 = docker_tester.containers.get("distsys-mq-worker-1")
    worker2 = docker_tester.containers.get("distsys-mq-worker-2")
    # Prevent the healthy worker from completing the whole batch before the fault.
    worker2.kill()
    subscribed = wait_for_worker_ready(worker1)
    worker1.pause()
    pending_ids = post_images(10)
    wait_for_pending_tasks(subscribed)
    worker1.kill()
    worker2.start()
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_two_faulty_workers(docker_tester):
    worker1 = docker_tester.containers.get("distsys-mq-worker-1")
    worker2 = docker_tester.containers.get("distsys-mq-worker-2")
    subscribed1 = wait_for_worker_ready(worker1)
    subscribed2 = wait_for_worker_ready(worker2)
    worker1.pause()
    worker2.pause()
    pending_ids = post_images(10)
    wait_for_pending_tasks(subscribed1 or subscribed2)
    worker1.kill()
    worker2.kill()
    worker1.start()
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_faulty_worker_and_rabbit_restart(docker_tester):
    worker1 = docker_tester.containers.get("distsys-mq-worker-1")
    worker2 = docker_tester.containers.get("distsys-mq-worker-2")
    rabbit = docker_tester.containers.get("distsys-mq-rabbitmq-1")
    # No worker may finish tasks before the broker loses its in-memory state.
    worker1.kill()
    worker2.kill()
    pending_ids = post_images(10)
    check_task_queue(10, 10)
    rabbit.kill()
    rabbit.start()
    wait_for_task_queue(lambda queue: True, "broker restart", max_attempts=20)
    worker2.start()
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker']])
def test_total_eclipse_of_the_heart(docker_tester):
    completed_ids = post_images(10)
    wait_and_check_results(completed_ids, 10, docker_client=docker_tester)

    worker1 = docker_tester.containers.get("distsys-mq-worker-1")
    worker2 = docker_tester.containers.get("distsys-mq-worker-2")
    rabbit = docker_tester.containers.get("distsys-mq-rabbitmq-1")
    worker1.kill()
    worker2.kill()
    rabbit.kill()

    # These URLs repeat the first batch, but every request needs a new ID.
    pending_ids = post_images(10, known_ids=completed_ids)
    # Also accept requests after the client has detected the broken connection.
    time.sleep(20)
    later_ids = post_images(10, known_ids=completed_ids | pending_ids)
    wait_and_check_results(completed_ids, 10, docker_client=docker_tester)

    rabbit.start()
    wait_for_task_queue(lambda queue: True, "broker restart", max_attempts=20)
    worker1.start()
    worker2.start()
    wait_and_check_results(completed_ids | pending_ids | later_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker', 'network-fault']])
def test_notifications_survive_rabbit_restart(docker_tester):
    worker1 = docker_tester.containers.get("distsys-mq-worker-1")
    worker2 = docker_tester.containers.get("distsys-mq-worker-2")
    rabbit = docker_tester.containers.get("distsys-mq-rabbitmq-1")
    fault = docker_tester.containers.get("distsys-mq-network-fault-1")
    worker1.kill()
    worker2.kill()
    pending_ids = post_images(10)
    check_task_queue(10, 10)

    rabbit.reload()
    address = next(network['IPAddress'] for network in
                   rabbit.attrs['NetworkSettings']['Networks'].values() if network['IPAddress'])
    with block_amqp_link(fault, address):
        worker1.start()
        worker2.start()
        wait_for_queued_notifications()
        # Notifications cannot reach the API; tasks may still await acknowledgement.
        wait_and_check_results(set(), 1, docker_client=docker_tester)
        worker1.kill()
        worker2.kill()
        rabbit.kill()
        rabbit.start()
        wait_for_task_queue(lambda queue: True, "broker restart", max_attempts=20)

    worker1.start()
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)


@pytest.mark.parametrize("services", [['rabbitmq', 'server', 'worker', 'network-fault']])
def test_duplicate_deliveries(docker_tester):
    worker1 = docker_tester.containers.get("distsys-mq-worker-1")
    worker2 = docker_tester.containers.get("distsys-mq-worker-2")
    server = docker_tester.containers.get("distsys-mq-server-1")
    rabbit = docker_tester.containers.get("distsys-mq-rabbitmq-1")
    fault = docker_tester.containers.get("distsys-mq-network-fault-1")
    worker1.kill()
    worker2.kill()
    pending_ids = post_images(10)
    check_task_queue(10, 10)
    task_copies = duplicate_queued_messages('task_queue')
    assert task_copies == 20
    check_task_queue(task_copies, 10)

    rabbit.reload()
    address = next(network['IPAddress'] for network in
                   rabbit.attrs['NetworkSettings']['Networks'].values() if network['IPAddress'])
    with block_amqp_link(fault, address):
        # Close the blocked connection so in-flight notifications become ready.
        disconnect_from_broker(server)
        worker1.start()
        worker2.start()
        queues = wait_for_queued_notifications(ready_only=True)
        wait_and_check_results(set(), 1, docker_client=docker_tester)
        worker1.kill()
        worker2.kill()
        notification_copies = {
            queue['name']: duplicate_queued_messages(queue['name'])
            for queue in queues if queue['name'] != 'task_queue' and queue.get('messages', 0) > 0
        }
        assert notification_copies
        wait_for_queue_counts(notification_copies)

    worker1.start()
    wait_and_check_results(pending_ids, 10, docker_client=docker_tester)
    # The first complete result list may precede processing of later copies.
    wait_for_queue_counts({name: 0 for name in notification_copies})
    wait_and_check_results(pending_ids, 1, docker_client=docker_tester)


# Utils ===============================================================================================================

@pytest.fixture
def docker_tester(services):
    print()
    try:
        run_docker_compose_up(services)
        check_server_endpoint()
        client = docker.from_env()
        yield client
    finally:
        print()
        run_docker_compose_down()


def run_docker_compose_up(services):
    command = ["docker", "compose", "--ansi", "never", "up", "--force-recreate"]
    for service in services:
        command.append(service)
    subprocess.Popen(command, cwd=Path(__file__).parent.parent.absolute(), stdout=None, stderr=None)


def run_docker_compose_down():
    command = ["docker", "compose", "down", "--volumes"]
    subprocess.run(command, cwd=Path(__file__).parent.parent.absolute(), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def check_server_endpoint(max_attempts=20):
    attempt = 0
    while True:
        attempt += 1
        try:
            logger.info(f"Sending GET {IMAGES_ENDPOINT}")
            response = requests.get(IMAGES_ENDPOINT, timeout=3)
            assert response.status_code == 200, response.text
            logger.info(f"Attempt {attempt} succeeded: got response, server endpoint is ready")
            return
        except Exception as e:
            logger.error(f"Attempt {attempt} failed: {e}")
        if attempt == max_attempts:
            logger.error(f"Max attempts reached, give up")
            pytest.fail("Server endpoint is not ready")
        logger.info(f"Retry in 3 seconds...")
        time.sleep(3)


def post_images(num_requests, known_ids=None):
    known_ids = set() if known_ids is None else known_ids
    pending_ids = set()
    for i in range(num_requests):
        input_data = {"image_url": f"https://somehost.com/some-image-{i}.jpg"}
        logger.info(f"Sending POST {IMAGES_ENDPOINT} {input_data}")
        try:
            response = requests.post(IMAGES_ENDPOINT, json=input_data, timeout=3)
        except Exception as e:
            logger.error(f"Request failed: {e}")
            pytest.fail(f"Failed to post image: {e}")
        logger.info(f"Got response: {response.status_code} {response.text.rstrip()}")
        assert response.status_code == 200
        assert 'image_id' in response.json()
        image_id = response.json()['image_id']
        assert type(image_id) in (str, int), 'image_id must be a string or integer'
        assert image_id not in known_ids, f'Reused image_id: {image_id}'
        assert image_id not in pending_ids
        pending_ids.add(image_id)
    return pending_ids


def wait_and_check_results(pending_ids, max_attempts, docker_client, check_captions=True):
    expected_count = len(pending_ids)
    attempt = 0
    while True:
        attempt += 1
        try:
            logger.info(f"Sending GET {IMAGES_ENDPOINT}")
            response = requests.get(IMAGES_ENDPOINT, timeout=3)
            logger.info(f"Got response: {response.status_code} {response.text.rstrip()}")
            assert response.status_code == 200
            assert 'image_ids' in response.json()
            ready_ids = response.json()['image_ids']
            assert isinstance(ready_ids, list), 'image_ids must be a list'
            assert all(type(image_id) in (str, int) for image_id in ready_ids)
            ready_set = set(ready_ids)
            assert len(ready_ids) == len(ready_set), 'image_ids contains duplicates'
            assert ready_set <= pending_ids, 'image_ids contains unexpected IDs'
            count = len(ready_set)
            if ready_set == pending_ids:
                if check_captions:
                    for image_id in pending_ids:
                        check_image_caption(image_id, docker_client)
                logger.info(f"Attempt {attempt} succeeded: got {count} results as expected")
                return
            else:
                logger.info(f"Attempt {attempt} not succeeded: got {count} results, expect {expected_count}")
        except Exception as e:
            logger.error(f"Attempt {attempt} failed: {e}")
        if attempt == max_attempts:
            logger.error(f"Max attempts reached, give up")
            pytest.fail("Server didn't return expected results")
        logger.info(f"Retry in 3 seconds...")
        time.sleep(3)


def check_image_caption(image_id, docker_client):
    logger.info(f"Sending GET {IMAGES_ENDPOINT}/{image_id}")
    try:
        response = requests.get(f'{IMAGES_ENDPOINT}/{image_id}', timeout=3)
    except Exception as e:
        logger.error(f"Request failed: {e}")
        pytest.fail(f"Failed to get image caption: {e}")
    logger.info(f"Got response: {response.status_code} {response.text.rstrip()}")
    assert response.status_code == 200
    assert 'caption' in response.json()
    assert isinstance(response.json()['caption'], str)
    saved_caption = read_saved_caption(image_id, docker_client)
    assert response.json()['caption'] == saved_caption, 'API caption differs from the saved result'


def read_saved_caption(image_id, docker_client):
    # Docker's archive API works even if the solution image has no shell or Python.
    server = docker_client.containers.get('distsys-mq-server-1')
    stream, _ = server.get_archive(f'/data/{image_id}.txt')
    with tarfile.open(fileobj=io.BytesIO(b''.join(stream))) as archive:
        files = [member for member in archive.getmembers() if member.isfile()]
        assert len(files) == 1, 'Expected one regular caption file'
        with archive.extractfile(files[0]) as result:
            return result.read().decode('utf-8')


def check_unprocessed_images(image_ids, docker_client):
    wait_and_check_results(set(), 1, docker_client)
    for image_id in image_ids:
        response = requests.get(f'{IMAGES_ENDPOINT}/{image_id}', timeout=3)
        assert response.status_code == 404, f'Result available without workers: {response.text}'


def wait_for_task_queue(predicate, description, max_attempts=10, interval=3):
    last_state = None
    for attempt in range(1, max_attempts + 1):
        try:
            response = requests.get(TASK_QUEUE_ENDPOINT, timeout=3)
            assert response.status_code == 200, response.text
            queue = response.json()
            last_state = {key: queue.get(key) for key in
                          ('messages_ready', 'messages_unacknowledged', 'consumers', 'message_stats')}
            if predicate(queue):
                logger.info(f"Task queue is ready: {description}")
                return queue
        except Exception as e:
            last_state = str(e)
        logger.info(f"Waiting for {description}: {last_state}")
        if attempt < max_attempts:
            time.sleep(interval)
    pytest.fail(f"Task queue did not reach {description}; last state: {last_state}")


def check_task_queue(expected_count, max_attempts):
    return wait_for_task_queue(lambda queue: queue['messages_ready'] == expected_count,
                               f'{expected_count} ready messages', max_attempts)


def wait_for_worker_ready(worker, max_attempts=10):
    last_state = None
    for attempt in range(1, max_attempts + 1):
        try:
            worker.reload()
            addresses = {network['IPAddress'] for network in
                         worker.attrs['NetworkSettings']['Networks'].values() if network['IPAddress']}
            response = requests.get(TASK_QUEUE_ENDPOINT, timeout=3)
            assert response.status_code == 200, response.text
            consumers = response.json().get('consumer_details', [])
            if any(consumer['channel_details']['peer_host'] in addresses for consumer in consumers):
                return True
            # Polling with basic.get is also valid; it has no registered consumer.
            response = requests.get(f'{BROKER_ENDPOINT}/channels', timeout=3)
            assert response.status_code == 200, response.text
            channels = response.json()
            last_state = [{'peer_host': channel['connection_details']['peer_host'],
                           'message_stats': channel.get('message_stats')} for channel in channels
                          if channel['connection_details']['peer_host'] in addresses]
            for channel in channels:
                if channel['connection_details']['peer_host'] in addresses:
                    stats = channel.get('message_stats') or {}
                    if any(stats.get(name, 0) > 0 for name in ('get', 'get_no_ack', 'get_empty')):
                        return False
        except Exception as e:
            last_state = str(e)
        if attempt < max_attempts:
            time.sleep(3)
    pytest.fail(f"Worker {worker.name} did not start consuming tasks; last state: {last_state}")


def wait_for_pending_tasks(subscribed):
    if subscribed:
        # Delivery statistics can arrive before the unacknowledged count.
        # Automatic ACK is deliberately allowed through this barrier: after the
        # kill, the missing results must expose the lost deliveries instead.
        predicate = lambda queue: (queue['messages_unacknowledged'] > 0 or
                                   any((queue.get('message_stats') or {}).get(name, 0) > 0
                                       for name in ('deliver', 'deliver_no_ack')))
        description = 'a task delivered to a paused worker'
    else:
        predicate = lambda queue: queue['messages_ready'] == 10
        description = '10 tasks waiting for a polling worker'
    # Statistics refresh every 1s in rabbitmq.conf.
    return wait_for_task_queue(predicate, description, max_attempts=10, interval=1)


def wait_for_queued_notifications(max_attempts=20, ready_only=False):
    last_state = None
    for attempt in range(1, max_attempts + 1):
        try:
            response = requests.get(f'{BROKER_ENDPOINT}/queues/%2F', timeout=3)
            assert response.status_code == 200, response.text
            queues = response.json()
            last_state = {queue['name']: queue.get('messages') for queue in queues}
            # Notifications may be batched or spread across multiple queues.
            # Workers may keep tasks unacknowledged until the API receives a result.
            notification_count = sum(count or 0 for name, count in last_state.items()
                                     if name != 'task_queue')
            ready = sum(queue.get('messages_ready', 0) or 0 for queue in queues
                        if queue['name'] != 'task_queue')
            if 'task_queue' in last_state and notification_count > 0 and (not ready_only or ready == notification_count):
                return queues
        except Exception as e:
            last_state = str(e)
        if attempt < max_attempts:
            time.sleep(1)
    pytest.fail(f'No queued notifications reached the required state; last state: {last_state}')


@contextmanager
def block_amqp_link(fault, address):
    rules = [('INPUT', '-s', address, '--sport'), ('OUTPUT', '-d', address, '--dport')]
    installed = []
    try:
        for chain, direction, host, port in rules:
            rule = [chain, direction, host, '-p', 'tcp', port, '5672', '-j', 'DROP']
            result = fault.exec_run(['iptables', '-I', *rule])
            assert result.exit_code == 0, result.output.decode(errors='replace')
            installed.append(rule)
        yield
    finally:
        for rule in reversed(installed):
            result = fault.exec_run(['iptables', '-D', *rule])
            assert result.exit_code == 0, result.output.decode(errors='replace')


def duplicate_queued_messages(queue_name):
    # Consumers/producers of this queue are stopped or disconnected by the test.
    messages = []
    while True:
        response = requests.post(
            f'{BROKER_ENDPOINT}/queues/%2F/{quote(queue_name, safe="")}/get',
            json={'count': 100, 'ackmode': 'ack_requeue_false', 'encoding': 'base64'},
            timeout=3)
        assert response.status_code == 200, response.text
        batch = response.json()
        if not batch:
            break
        messages.extend(batch)
    assert messages, f'No messages to duplicate in {queue_name}'
    for message in messages:
        for _ in range(2):
            response = requests.post(f'{BROKER_ENDPOINT}/exchanges/%2F/amq.default/publish',
                json={'routing_key': queue_name, 'properties': message['properties'],
                      'payload': message['payload'], 'payload_encoding': 'base64'}, timeout=3)
            assert response.status_code == 200, response.text
            assert response.json()['routed'] is True, f'Copy was not routed to {queue_name}'
    logger.info(f'Queued {len(messages) * 2} copies in {queue_name}')
    return len(messages) * 2


def disconnect_from_broker(container):
    container.reload()
    addresses = {network['IPAddress'] for network in
                 container.attrs['NetworkSettings']['Networks'].values() if network['IPAddress']}
    response = requests.get(f'{BROKER_ENDPOINT}/connections', timeout=3)
    assert response.status_code == 200, response.text
    for connection in response.json():
        if connection['peer_host'] in addresses:
            response = requests.delete(
                f'{BROKER_ENDPOINT}/connections/{quote(connection["name"], safe="")}', timeout=3)
            assert response.status_code in (204, 404), response.text


def wait_for_queue_counts(expected, max_attempts=30):
    last_state = None
    for attempt in range(1, max_attempts + 1):
        try:
            last_state = {}
            for name in expected:
                response = requests.get(f'{BROKER_ENDPOINT}/queues/%2F/{quote(name, safe="")}', timeout=3)
                assert response.status_code == 200, response.text
                last_state[name] = response.json().get('messages')
            if last_state == expected:
                return
        except Exception as e:
            last_state = str(e)
        if attempt < max_attempts:
            time.sleep(1)
    pytest.fail(f'Queues did not reach {expected}; last state: {last_state}')