import time
from typing import Dict, Optional
import forcen_public_utils.checked as ch


class MultiTimestampSync:
    """Modifies timestamps from multiple clock sources
    to reference a clientside reset event. Assumes clocks
    have the same scale, different phase"""

    def __init__(self):
        self.offsets: Dict[int, float] = dict()
        self.reset_timestamp = time.time()

    def reset(self, reset_timestamp=None):
        """server_timestamp, client_timestamp are the timestamps corresponding to
        a single network transmission"""
        if reset_timestamp is None:
            reset_timestamp = time.time()

        self.reset_timestamp = reset_timestamp
        self.offsets = dict()

    def generate_from_server(
        self, source: int, server_timestamp: float, client_timestamp: float
    ) -> float:
        if source in self.offsets.keys():
            return server_timestamp + self.offsets[source]

        self.offsets[source] = (
            client_timestamp - server_timestamp - self.reset_timestamp
        )
        return client_timestamp - self.reset_timestamp

    def generate_from_client(self, client_timestamp: Optional[float] = None) -> float:
        if client_timestamp is None:
            client_timestamp = time.time()

        return client_timestamp - self.reset_timestamp
