import time
from rich.pretty import pprint
from typing import Iterable
from opentelemetry._metrics.measurement import Measurement
from opentelemetry.sdk._metrics import MeterProvider
from opentelemetry.sdk._metrics.metric_reader import MetricReader
from opentelemetry.sdk._metrics.point import AggregationTemporality, Metric


class MyReader(MetricReader):
    def _receive_metrics(self, metrics: Iterable[Metric]) -> None:
        """Called by `MetricReader.collect` when it receives a batch of metrics"""
        pprint(metrics)

    def shutdown(self) -> bool:
        print("shutdown called")
        return True


def cpu_time_callback() -> Iterable[Measurement]:
    with open("/proc/stat") as procstat:
        procstat.readline()  # skip the first line
        for line in procstat:
            if not line.startswith("cpu"):
                break
            cpu, *states = line.split()
            yield Measurement(
                int(states[0]) // 100, {"cpu": cpu, "state": "user"}
            )
            yield Measurement(
                int(states[1]) // 100, {"cpu": cpu, "state": "nice"}
            )
            # ... other states


def main() -> None:
    reader = MyReader(preferred_temporality=AggregationTemporality.CUMULATIVE)
    mp = MeterProvider(metric_readers=[reader])
    meter = mp.get_meter("hellometer")
    mycounter = meter.create_counter("mycounter")
    meter.create_observable_counter(
        "system.cpu.time",
        callback=cpu_time_callback,
        unit="ms",
        description="System's CPU time",
    )
    for _ in range(20):
        mycounter.add(12, {"path": "/hello"})
        mycounter.add(10, {"path": "/world"})

    reader.collect()

    time.sleep(2)
    mp.shutdown()


if __name__ == "__main__":
    main()