"""Clock synchronization helpers.

This module provides ``ClockSynchronizer``, a small utility that maps a
remote 32-bit wrapping timetag into the local monotonic clock. It uses a
linear regression model (NumPy) to estimate offset and drift and estimates
one-way latency.

Note: NumPy is required by this module; an ImportError is raised when it is
not available.
"""

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:
    """Synchronize a remote 32-bit wrapping counter to local monotonic time.

    The synchronizer records pairs of (server_timetag, client_reception_time),
    unwraps the server counter across wrap-around events and fits a linear
    model (client_time = m * server_ticks + b). The fit is performed using
    NumPy's polyfit when a configurable minimum number of samples is
    available.
    """

    # 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):
        """Create a new ClockSynchronizer.

        Args:
            history_size: Maximum number of recent samples retained for fitting.
            min_samples_for_fit: Minimum samples required before performing a
                regression fit.
        Raises:
            ImportError: If NumPy is not available on the system.
        """
        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):
        """Add a new (server_timetag, client_time) sample.

        The method will unwrap the 32-bit server timetag accounting for wrap
        events and append the unwrapped pair to the internal history. If enough
        samples are present a new regression fit will be computed.

        Args:
            raw_server_timetag: Raw 32-bit server timetag.
            client_reception_time: Local monotonic time of the packet 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):
        """Internal: fit a linear model to the sample history.

        This updates slope/intercept (m, b) and computes an estimated average
        one-way latency. Must be called while holding the instance lock.
        """
        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:
        """Map a raw server timetag to an estimated client monotonic time.

        The method unwraps the provided 32-bit timetag using the current
        wrap count and applies the linear model (m, b) to estimate the
        corresponding local monotonic time.

        Args:
            raw_server_timetag: Raw 32-bit server timetag.

        Returns:
            Estimated client monotonic time (float).
        """
        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:
        """Return the estimated average one-way server->client latency in seconds.

        Returns 0.0 when insufficient data is available to compute the metric.
        """
        with self._lock:
            return self._average_latency_s
