from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import Protocol

from observability import elapsed_ms, timer_start


@dataclass(frozen=True)
class SinkSample:
    decode_duration_ms: float
    write_duration_ms: float


class Sink[InputT](Protocol):
    def write(self, value: InputT) -> None: ...


class SampledSink[InputT](Sink[InputT], Protocol):
    def write_sample(self, value: InputT) -> SinkSample: ...


@dataclass(frozen=True)
class DecodedSink[InputT, DecodedT]:
    decode: Callable[[InputT], DecodedT]
    write_decoded: Callable[[DecodedT], None]

    def write(self, value: InputT) -> None:
        self.write_decoded(self.decode(value))

    def write_sample(self, value: InputT) -> SinkSample:
        decode_start = timer_start()
        decoded = self.decode(value)
        decode_duration_ms = elapsed_ms(decode_start)

        write_start = timer_start()
        self.write_decoded(decoded)
        write_duration_ms = elapsed_ms(write_start)
        return SinkSample(decode_duration_ms=decode_duration_ms, write_duration_ms=write_duration_ms)


@dataclass(frozen=True)
class RawSink[InputT]:
    write_raw: Callable[[InputT], None]

    def write(self, value: InputT) -> None:
        self.write_raw(value)

    def write_sample(self, value: InputT) -> SinkSample:
        write_start = timer_start()
        self.write_raw(value)
        return SinkSample(decode_duration_ms=0.0, write_duration_ms=elapsed_ms(write_start))


__all__ = ["DecodedSink", "RawSink", "SampledSink", "Sink", "SinkSample"]