"""
Provides a ClockSynchronizer class to model the relationship between a remote
server's wrapping 32-bit timetag and the local monotonic clock.
"""

import collections
import threading
import time
from typing import List, Tuple

# NumPy is a strong recommendation for linear regression.
# If it's not already a dependency, it should be added.
try:
    import numpy as np

    NUMPY_AVAILABLE = True
except ImportError:
    NUMPY_AVAILABLE = False


class ClockSynchronizer:
    """
    Synchronizes a remote wrapping 32-bit counter with the local monotonic clock
    using linear regression to model clock offset and drift.
    """

    # Constants for a 32-bit counter
    _COUNTER_MAX = 2**32
    _WRAP_THRESHOLD = 2**31  # Detect wrap if decrease is > half the max value

    def __init__(self, history_size: int = 100, min_samples_for_fit: int = 10):
        """
        Initializes the ClockSynchronizer.

        Args:
            history_size: The number of recent samples to use for regression.
            min_samples_for_fit: The minimum number of samples required to
                                perform a linear regression fit.
        """
        if not NUMPY_AVAILABLE:
            raise ImportError("NumPy is required for the ClockSynchronizer.")

        self._lock = threading.Lock()
        self._history: collections.deque = collections.deque(maxlen=history_size)
        self._min_samples = min_samples_for_fit

        # State for timestamp unwrapping
        self._wrap_count: int = 0
        self._last_raw_timetag: int | None = None

        # Linear model parameters: client_time = m * server_unwrapped_ticks + b
        self._m: float = 0.0  # Slope (client seconds per server tick)
        self._b: float = 0.0  # Intercept (client time when server time was 0)

        # Estimated one-way latency from server to client
        self._average_latency_s: float = 0.0

    def add_sample(self, raw_server_timetag: int, client_reception_time: float):
        """
        Adds a new sample pair to update the synchronization model.

        Args:
            raw_server_timetag: The raw 32-bit timetag from the server.
            client_reception_time: The local monotonic time of reception.
        """
        with self._lock:
            # --- Timestamp Unwrapping Logic ---
            if self._last_raw_timetag is None:
                # First sample, assume no wraps yet.
                self._last_raw_timetag = raw_server_timetag
            else:
                # Check for a wrap-around
                diff = self._last_raw_timetag - raw_server_timetag
                if diff > self._WRAP_THRESHOLD:
                    self._wrap_count += 1

            self._last_raw_timetag = raw_server_timetag
            unwrapped_timetag = (
                raw_server_timetag + self._wrap_count * self._COUNTER_MAX
            )

            # Add the new sample to history and update the model
            self._history.append((unwrapped_timetag, client_reception_time))
            self._update_model()

    def _update_model(self):
        """
        Performs linear regression on the stored history to update the
        model parameters (m and b) and the average latency.
        This method must be called within a locked context.
        """
        if len(self._history) < self._min_samples:
            # Not enough data for a reliable fit
            return

        x_vals = np.array([sample[0] for sample in self._history])
        y_vals = np.array([sample[1] for sample in self._history])

        try:
            m, b = np.polyfit(x_vals, y_vals, 1)
            self._m = m
            self._b = b

            # --- Calculate Average Latency ---
            # Estimated generation time for each sample based on the new model
            estimated_generation_times = self._m * x_vals + self._b
            # Latency is the difference between reception and estimated generation
            latencies = y_vals - estimated_generation_times
            # Update the average latency, filtering out negative values which are artifacts
            positive_latencies = latencies[latencies >= 0]
            if len(positive_latencies) > 0:
                self._average_latency_s = np.mean(positive_latencies)
            else:
                self._average_latency_s = 0.0

        except np.linalg.LinAlgError:
            pass

    def to_client_time(self, raw_server_timetag: int) -> float:
        """
        Estimates the equivalent local client monotonic time for a given raw
        server timetag.

        Args:
            raw_server_timetag: The raw 32-bit timetag from the server.

        Returns:
            The estimated client monotonic time when the event occurred.
        """
        with self._lock:
            current_wrap_count = self._wrap_count
            if self._last_raw_timetag is not None:
                diff = self._last_raw_timetag - raw_server_timetag
                if diff < -self._WRAP_THRESHOLD:
                    current_wrap_count -= 1

            unwrapped_timetag = (
                raw_server_timetag + current_wrap_count * self._COUNTER_MAX
            )

            estimated_time = self._m * unwrapped_timetag + self._b
            return estimated_time

    def get_average_latency_s(self) -> float:
        """
        Returns the current estimated average one-way network latency from
        server to client in seconds.

        Returns:
            The average latency in seconds, or 0.0 if not yet computed.
        """
        with self._lock:
            return self._average_latency_s