import copy
import json
import os
import pathlib
import queue
import re
import socket
import subprocess
import threading
import time
from datetime import datetime, timezone
from typing import Dict

import pytest

test_message = {'author': 'alice', 'text': 'hello'}

PROTO_DIR = pathlib.Path(__file__).resolve().parent.parent / 'solution' / 'proto'

SOCKET_CONNECT_TIMEOUT_S = 1
SERVICE_READY_TIMEOUT_S = 20
SERVICE_RETRY_INTERVAL_S = 0.5
GRPC_CALL_TIMEOUT_S = 5
GRPC_PROCESS_TIMEOUT_S = 10
GRPC_STREAM_TIMEOUT_S = 60
PROCESS_STOP_TIMEOUT_S = 5
MESSAGE_TIMEOUT_S = 10


def wait_for_socket(host, port):
    deadline = time.monotonic() + SERVICE_READY_TIMEOUT_S
    last_exception = None
    while True:
        try:
            with socket.create_connection((host, port), timeout=SOCKET_CONNECT_TIMEOUT_S):
                pass
            return
        except OSError as exc:
            last_exception = exc

        remaining = deadline - time.monotonic()
        if remaining <= 0:
            pytest.fail(
                f'timed out waiting for TCP endpoint {host}:{port}: {last_exception}',
                pytrace=False,
            )
        time.sleep(min(SERVICE_RETRY_INTERVAL_S, remaining))


@pytest.fixture(scope='session')
def server_addr():
    addr = os.environ.get('MESSENGER_TEST_SERVER_ADDR', '127.0.0.1:51075')
    host = addr.split(':')[0]
    port = int(addr.split(':')[1])
    wait_for_socket(host, port)
    yield addr


def send_message(server_address, message: Dict[str, str]) -> Dict[str, str]:
    grpcurl_cmd = ['grpcurl',
                   '-max-time', str(GRPC_CALL_TIMEOUT_S),
                   '-import-path', str(PROTO_DIR),
                   '-proto', 'messenger.proto',
                   '-d',
                   json.dumps(message),
                   '-plaintext',
                   server_address,
                   'mes_grpc.MessengerServer/SendMessage']
    try:
        completed = subprocess.run(
            grpcurl_cmd,
            capture_output=True,
            check=False,
            timeout=GRPC_PROCESS_TIMEOUT_S,
        )
    except subprocess.TimeoutExpired:
        pytest.fail(
            f'grpcurl did not finish within {GRPC_PROCESS_TIMEOUT_S} seconds',
            pytrace=False,
        )
    assert completed.returncode == 0, completed.stderr
    assert len(completed.stderr) == 0, completed.stderr
    output_str = completed.stdout.decode('ascii')
    output = json.loads(output_str)

    message_with_timestamp = copy.deepcopy(message)
    message_with_timestamp['sendTime'] = output['sendTime']
    return message_with_timestamp


class MessageStream:
    def __init__(self, server_address):
        grpcurl_cmd = ['grpcurl',
                       '-max-time', str(GRPC_STREAM_TIMEOUT_S),
                       '-import-path', str(PROTO_DIR),
                       '-proto', 'messenger.proto',
                       '-plaintext',
                       server_address,
                       'mes_grpc.MessengerServer/ReadMessages']
        self._process = subprocess.Popen(
            grpcurl_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        self._messages = queue.Queue()
        self._reader_error = None
        self._reader_finished = threading.Event()
        self._closed = False
        self._reader = threading.Thread(target=self._read_messages, daemon=True)
        self._reader.start()

    def _read_messages(self):
        try:
            message_lines = []
            for line in self._process.stdout:
                message_lines.append(line)
                if line.rstrip() == '}':
                    self._messages.put(json.loads(''.join(message_lines)))
                    message_lines = []
        except Exception as exc:
            self._reader_error = exc
        finally:
            self._reader_finished.set()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def _failure_detail(self):
        if self._reader_error is not None:
            return f'reader failed: {self._reader_error}'
        returncode = self._process.poll()
        if returncode is not None:
            stderr = self._process.stderr.read().strip()
            return f'grpcurl exited with status {returncode}: {stderr}'
        return 'grpcurl is still running but produced no matching message'

    def _get_message(self, deadline, timeout_message):
        while True:
            if self._reader_error is not None:
                raise AssertionError(self._failure_detail())
            if self._reader_finished.is_set() and self._messages.empty():
                raise AssertionError(self._failure_detail())

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise AssertionError(f'{timeout_message}: {self._failure_detail()}')
            try:
                return self._messages.get(timeout=min(remaining, 0.1))
            except queue.Empty:
                pass

    def wait_for_message(self, expected_message, timeout, preceding_messages=None):
        deadline = time.monotonic() + timeout
        while True:
            message = self._get_message(
                deadline,
                f'timed out waiting for message {expected_message}',
            )
            if message == expected_message:
                return
            if preceding_messages is not None:
                preceding_messages.append(message)

    def read_messages(self, count, timeout):
        deadline = time.monotonic() + timeout
        messages = []
        while len(messages) < count:
            messages.append(self._get_message(
                deadline,
                f'timed out after receiving {len(messages)} of {count} messages',
            ))
        return messages

    def close(self):
        if self._closed:
            return
        self._closed = True
        if self._process.poll() is None:
            self._process.terminate()
        try:
            self._process.wait(timeout=PROCESS_STOP_TIMEOUT_S)
        except subprocess.TimeoutExpired:
            self._process.kill()
            try:
                self._process.wait(timeout=PROCESS_STOP_TIMEOUT_S)
            except subprocess.TimeoutExpired:
                pytest.fail('grpcurl did not exit after SIGKILL', pytrace=False)
        self._reader.join(timeout=PROCESS_STOP_TIMEOUT_S)
        assert not self._reader.is_alive()


def wait_for_streams(server_address, streams):
    preceding_messages = [[] for _ in streams]
    probes = []
    for attempt in range(10):
        probe = send_message(
            server_address,
            {'author': 'StreamProbe', 'text': f'probe #{attempt}'},
        )
        probes.append(probe)
        streams_ready = True
        for index, stream in enumerate(streams):
            try:
                stream.wait_for_message(
                    probe,
                    timeout=1,
                    preceding_messages=preceding_messages[index],
                )
            except AssertionError:
                streams_ready = False
        if streams_ready:
            return [
                [message for message in messages if message not in probes]
                for messages in preceding_messages
            ]
    raise AssertionError('ReadMessages streams did not become ready')


def timestamp_key(timestamp):
    match = re.fullmatch(r'(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(?:\.(\d{1,9}))?Z', timestamp)
    assert match is not None, f'invalid protobuf timestamp: {timestamp}'
    seconds = int(datetime.strptime(match.group(1), '%Y-%m-%dT%H:%M:%S')
                  .replace(tzinfo=timezone.utc).timestamp())
    nanos = int((match.group(2) or '').ljust(9, '0'))
    return seconds, nanos


def test_send_smoke(server_addr):
    send_message(server_addr, test_message)


def test_send_returns_ascending_time(server_addr):
    outputs = []
    for _ in range(10):
        outputs.append(send_message(server_addr, test_message))

    for output1, output2 in zip(outputs, outputs[1:]):
        assert timestamp_key(output1['sendTime']) < timestamp_key(output2['sendTime'])


def test_get_messages_smoke(server_addr):
    with MessageStream(server_addr) as stream:
        wait_for_streams(server_addr, [stream])
        test_message_with_timestamp = send_message(server_addr, test_message)
        messages = stream.read_messages(1, timeout=MESSAGE_TIMEOUT_S)

    assert len(messages) == 1
    assert messages[0] == test_message_with_timestamp


def test_get_only_sends_new(server_addr):
    messages1 = []
    messages3 = []
    n1, n2, n3 = 2, 3, 4

    with MessageStream(server_addr) as stream:
        messages_before_ready = wait_for_streams(server_addr, [stream])
        assert messages_before_ready == [[]]
        for _ in range(n1):
            message = send_message(server_addr, test_message)
            messages1.append(message)
        messages = stream.read_messages(n1, timeout=MESSAGE_TIMEOUT_S)

    assert len(messages1) == len(messages)
    for m1, m2 in zip(messages1, messages):
        assert m1 == m2

    for _ in range(n2):
        send_message(server_addr, test_message)

    with MessageStream(server_addr) as stream:
        messages_before_ready = wait_for_streams(server_addr, [stream])
        assert messages_before_ready == [[]]
        for _ in range(n3):
            message = send_message(server_addr, test_message)
            messages3.append(message)
        messages = stream.read_messages(n3, timeout=MESSAGE_TIMEOUT_S)

    assert len(messages3) == len(messages)
    for m1, m2 in zip(messages3, messages):
        assert m1 == m2