import time
from typing import Callable, List

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import StatusCode


class SpanExporterSpy(InMemorySpanExporter):
    def get_published_span_by_name(self, name: str) -> ReadableSpan:
        return self.get_published_span(lambda span: span.name == name)

    def get_published_span(self, func: Callable[[ReadableSpan], bool]) -> ReadableSpan:
        found = self.get_published_spans(func=func)

        if len(found) == 0:
            raise RuntimeError("No span was published")

        if len(found) > 1:
            raise RuntimeError("More than one span was published")

        return found[0]

    def get_published_spans(self, func: Callable[[ReadableSpan], bool]) -> List[ReadableSpan]:
        return list(filter(func, self.get_finished_spans()))

    def wait_for_span_with_name(self, name: str, timeout: float = 5) -> ReadableSpan:
        return self.wait_for_span(
            func=lambda span: span.name == name,
            timeout=timeout
        )

    def wait_for_successful_span_with_name(self, name: str, timeout: float = 5) -> ReadableSpan:
        return self.wait_for_span(
            func=lambda span: span.status.status_code == StatusCode.OK and span.name == name,
            timeout=timeout
        )

    def wait_for_span_with_name_and_attribute(self, name, attr_key: str, attr_value, timeout: float = 5):
        def my_filter(span):
            return span.name == name and attr_key in span.attributes and span.attributes[attr_key] == attr_value

        return self.wait_for_span(
            func=my_filter,
            timeout=timeout
        )

    def wait_for_span(self, func: Callable[[ReadableSpan], bool], timeout: float = 5) -> ReadableSpan:
        end_time = time.time() + timeout
        while time.time() < end_time:
            received = list(filter(func, self.get_finished_spans()))
            if len(received) != 0:
                return received[0]
            time.sleep(0)

        raise RuntimeError("Span was never published")
