import math
import random
import logging
import threading
import time
from contextlib import contextmanager
from dataclasses import dataclass
from functools import wraps
from typing import Callable, Dict, List, Union

Number = Union[int, float]

TICK_INTERVAL = 5.0


class EWMA:
    """
    Exponentially-weighted moving average, modeled after Dropwizard / Coda Hale Metrics.
    Rates are expressed in events / second.
    """

    M1_ALPHA = 1 - math.exp(-TICK_INTERVAL / 60.0 / 1.0)
    M5_ALPHA = 1 - math.exp(-TICK_INTERVAL / 60.0 / 5.0)
    M15_ALPHA = 1 - math.exp(-TICK_INTERVAL / 60.0 / 15.0)

    def __init__(self, alpha: float, interval: float = TICK_INTERVAL) -> None:
        self._alpha = alpha
        self._interval = interval
        self._uncounted = 0
        self._rate = 0.0
        self._initialized = False
        self._lock = threading.Lock()

    @classmethod
    def one_minute_ewma(cls) -> "EWMA":
        return cls(cls.M1_ALPHA)

    @classmethod
    def five_minute_ewma(cls) -> "EWMA":
        return cls(cls.M5_ALPHA)

    @classmethod
    def fifteen_minute_ewma(cls) -> "EWMA":
        return cls(cls.M15_ALPHA)

    def update(self, n: int) -> None:
        if n < 0:
            raise ValueError("n must be >= 0")
        with self._lock:
            self._uncounted += n

    def tick(self) -> None:
        with self._lock:
            count = self._uncounted
            self._uncounted = 0

            instant_rate = count / self._interval
            if self._initialized:
                self._rate += self._alpha * (instant_rate - self._rate)
            else:
                self._rate = instant_rate
                self._initialized = True

    def get_rate(self) -> float:
        with self._lock:
            return self._rate


class Counter:
    def __init__(self) -> None:
        self._value = 0
        self._lock = threading.Lock()

    def inc(self, amount: int = 1) -> None:
        if amount < 0:
            raise ValueError("amount must be >= 0")
        with self._lock:
            self._value += amount

    def dec(self, amount: int = 1) -> None:
        if amount < 0:
            raise ValueError("amount must be >= 0")
        with self._lock:
            self._value -= amount

    def get(self) -> int:
        with self._lock:
            return self._value

    def snapshot(self) -> Dict[str, Number]:
        return {"count": self.get()}


class Gauge:
    def __init__(self, func: Callable[[], Number]) -> None:
        self._func = func

    def get(self) -> Number:
        return self._func()

    def snapshot(self) -> Dict[str, Number]:
        return {"value": self.get()}


class Meter:
    """Dropwizard-like meter with (count and rates)"""

    def __init__(self, clock: Callable[[], float] = time.monotonic) -> None:
        self._clock = clock
        self._start_time = clock()
        self._count = 0
        self._lock = threading.Lock()

        self._m1 = EWMA.one_minute_ewma()
        self._m5 = EWMA.five_minute_ewma()
        self._m15 = EWMA.fifteen_minute_ewma()
        self._last_tick = self._start_time

    def _tick_if_necessary(self) -> None:
        old_tick = self._last_tick
        new_tick = self._clock()
        age = new_tick - old_tick

        if age < TICK_INTERVAL:
            return

        with self._lock:
            old_tick = self._last_tick
            new_tick = self._clock()
            age = new_tick - old_tick
            if age < TICK_INTERVAL:
                return

            new_interval_start = new_tick - (age % TICK_INTERVAL)
            required_ticks = int((new_interval_start - old_tick) / TICK_INTERVAL)
            self._last_tick = new_interval_start

        for _ in range(required_ticks):
            self._m1.tick()
            self._m5.tick()
            self._m15.tick()

    def mark(self, n: int = 1) -> None:
        if n < 0:
            raise ValueError("n must be >= 0")

        self._tick_if_necessary()
        with self._lock:
            self._count += n

        self._m1.update(n)
        self._m5.update(n)
        self._m15.update(n)

    def get_count(self) -> int:
        with self._lock:
            return self._count

    def get_mean_rate(self) -> float:
        with self._lock:
            count = self._count
            start_time = self._start_time

        elapsed = self._clock() - start_time
        if elapsed <= 0:
            return 0.0
        return count / elapsed

    def get_m1_rate(self) -> float:
        self._tick_if_necessary()
        return self._m1.get_rate()

    def get_m5_rate(self) -> float:
        self._tick_if_necessary()
        return self._m5.get_rate()

    def get_m15_rate(self) -> float:
        self._tick_if_necessary()
        return self._m15.get_rate()

    def snapshot(self) -> Dict[str, Number]:
        return {
            "count": self.get_count(),
            "mean_rate": self.get_mean_rate(),
            "m1_rate": self.get_m1_rate(),
            "m5_rate": self.get_m5_rate(),
            "m15_rate": self.get_m15_rate(),
        }


class Reservoir:
    """
    Fixed-size reservoir sampler for approximate quantiles with bounded memory.
    """

    def __init__(self, size: int = 1028, rand: Callable[[int, int], int] = random.randint) -> None:
        if size <= 0:
            raise ValueError("size must be > 0")
        self._size = size
        self._rand = rand
        self._values: List[float] = []
        self._count = 0
        self._lock = threading.Lock()

    def update(self, value: float) -> None:
        with self._lock:
            self._count += 1
            count = self._count

            if len(self._values) < self._size:
                self._values.append(value)
                return

            idx = self._rand(0, count - 1)
            if idx < self._size:
                self._values[idx] = value

    def snapshot(self) -> List[float]:
        with self._lock:
            return list(self._values)


class Timer:
    """
    Timer = histogram of durations + meter of event rates.
    Durations are in seconds.
    """

    def __init__(
        self,
        clock: Callable[[], float] = time.monotonic,
        reservoir_size: int = 1028,
    ) -> None:
        self._clock = clock
        self._reservoir = Reservoir(size=reservoir_size)
        self._meter = Meter(clock=clock)

        self._count = 0
        self._sum = 0.0
        self._min = math.inf
        self._max = 0.0
        self._lock = threading.Lock()

    def update(self, duration_seconds: float) -> None:
        if duration_seconds < 0:
            raise ValueError("duration_seconds must be >= 0")

        self._reservoir.update(duration_seconds)
        self._meter.mark()

        with self._lock:
            self._count += 1
            self._sum += duration_seconds
            if duration_seconds < self._min:
                self._min = duration_seconds
            if duration_seconds > self._max:
                self._max = duration_seconds

    @contextmanager
    def time(self):
        start = self._clock()
        try:
            yield
        finally:
            self.update(self._clock() - start)

    def record_callable(self, func: Callable, *args, **kwargs):
        start = self._clock()
        try:
            return func(*args, **kwargs)
        finally:
            self.update(self._clock() - start)

    def get_count(self) -> int:
        with self._lock:
            return self._count

    @staticmethod
    def _percentile(sorted_values: List[float], q: float) -> float:
        if not sorted_values:
            return 0.0
        if len(sorted_values) == 1:
            return sorted_values[0]

        pos = q * (len(sorted_values) - 1)
        lower = math.floor(pos)
        upper = math.ceil(pos)

        if lower == upper:
            return sorted_values[int(pos)]

        lower_value = sorted_values[lower]
        upper_value = sorted_values[upper]
        fraction = pos - lower
        return lower_value + (upper_value - lower_value) * fraction

    def snapshot(self) -> Dict[str, Number]:
        with self._lock:
            count = self._count
            total = self._sum
            min_value = self._min
            max_value = self._max

        if count == 0:
            return {
                "count": 0,
                "avg": 0.0,
                "min": 0.0,
                "max": 0.0,
                "p50": 0.0,
                "p95": 0.0,
                "p99": 0.0,
                "mean_rate": self._meter.get_mean_rate(),
                "m1_rate": self._meter.get_m1_rate(),
                "m5_rate": self._meter.get_m5_rate(),
                "m15_rate": self._meter.get_m15_rate(),
            }

        sample = self._reservoir.snapshot()
        sample.sort()

        return {
            "count": count,
            "avg": total / count,
            "min": min_value,
            "max": max_value,
            "p50": self._percentile(sample, 0.50),
            "p95": self._percentile(sample, 0.95),
            "p99": self._percentile(sample, 0.99),
            "mean_rate": self._meter.get_mean_rate(),
            "m1_rate": self._meter.get_m1_rate(),
            "m5_rate": self._meter.get_m5_rate(),
            "m15_rate": self._meter.get_m15_rate(),
        }


@dataclass
class _MetricEntry:
    metric_type: str
    metric: object


class MetricsRegistry:
    def __init__(self) -> None:
        self._metrics: Dict[str, _MetricEntry] = {}
        self._lock = threading.Lock()

    def counter(self, name: str) -> Counter:
        return self._get_or_create(name, "counter", Counter)

    def meter(self, name: str) -> Meter:
        return self._get_or_create(name, "meter", Meter)

    def timer(self, name: str, reservoir_size: int = 1028) -> Timer:
        with self._lock:
            existing = self._metrics.get(name)
            if existing is not None:
                if existing.metric_type != "timer":
                    raise ValueError(f"Metric {name!r} already exists as {existing.metric_type}")
                return existing.metric  # type: ignore[return-value]

            metric = Timer(reservoir_size=reservoir_size)
            self._metrics[name] = _MetricEntry("timer", metric)
            return metric

    def gauge(self, name: str, func: Callable[[], Number]) -> Gauge:
        with self._lock:
            existing = self._metrics.get(name)
            if existing is not None:
                if existing.metric_type != "gauge":
                    raise ValueError(f"Metric {name!r} already exists as {existing.metric_type}")
                return existing.metric  # type: ignore[return-value]

            metric = Gauge(func)
            self._metrics[name] = _MetricEntry("gauge", metric)
            return metric

    def _get_or_create(self, name: str, metric_type: str, cls):
        with self._lock:
            existing = self._metrics.get(name)
            if existing is not None:
                if existing.metric_type != metric_type:
                    raise ValueError(f"Metric {name!r} already exists as {existing.metric_type}")
                return existing.metric

            metric = cls()
            self._metrics[name] = _MetricEntry(metric_type, metric)
            return metric

    def get_metric(self, name: str):
        with self._lock:
            entry = self._metrics.get(name)
            return entry.metric if entry else None

    def to_dict(self) -> Dict[str, Dict[str, Number]]:
        with self._lock:
            items = list(self._metrics.items())

        result: Dict[str, Dict[str, Number]] = {}
        for name, entry in items:
            result[name] = entry.metric.snapshot()
        return result

    def to_flat_dict(self) -> Dict[str, Number]:
        nested = self.to_dict()
        flat: Dict[str, Number] = {}
        for metric_name, fields in nested.items():
            for field_name, value in fields.items():
                flat[f"{metric_name}.{field_name}"] = value
        return flat

    @staticmethod
    def _format_value(value: Number) -> str:
        if isinstance(value, float):
            rendered = f"{value:.6f}".rstrip("0").rstrip(".")
            if "." not in rendered:
                rendered += ".0"
            return rendered
        return str(value)

    def to_logfmt(self) -> str:
        flat = self.to_flat_dict()
        return " ".join(f"{key}={self._format_value(flat[key])}" for key in sorted(flat.keys()))


GLOBAL_REGISTRY = MetricsRegistry()


def start_metric_reporter(report_every_s=300.0):
    def _reporter():
        while True:
            time.sleep(report_every_s)
            try:
                logging.info("Metrics: %s", GLOBAL_REGISTRY.to_logfmt())
            except Exception:
                logging.exception("Metrics reporter failed")

    threading.Thread(target=_reporter, name="metrics-reporter", daemon=True).start()
