import time
from typing import Dict

from opentelemetry import metrics
from opentelemetry.exporter.otlp.proto.common.metrics_encoder import encode_metrics
from opentelemetry.proto.metrics.v1.metrics_pb2 import AggregationTemporality
from opentelemetry.sdk.metrics.export import MetricExporter, MetricExportResult, MetricsData, PeriodicExportingMetricReader
from opentelemetry.sdk.metrics.view import Aggregation
from opentelemetry.sdk.metrics import MeterProvider

from spotflow_device import StreamSender, DeviceClient

class SpotflowMetricExporter(MetricExporter):
    def __init__(
        self,
        stream_sender: StreamSender,
        preferred_temporality: Dict[type, AggregationTemporality] = None,
        preferred_aggregation: Dict[type, Aggregation] = None,
    ):
        super().__init__(
            preferred_temporality=preferred_temporality,
            preferred_aggregation=preferred_aggregation,
        )
        self.__stream_sender = stream_sender

    def export(self, metrics_data: MetricsData, timeout_millis: float = 10_000, **kwargs) -> MetricExportResult:
        print("Sending metrics.")
        # Serialize metrics data and send it via StreamSender
        self.__stream_sender.send_message(encode_metrics(metrics_data).SerializeToString(), batch_slice_id="metrics")
        return MetricExportResult.SUCCESS

    def force_flush(self, timeout_millis: float = 10_000) -> bool:
        return True

    def shutdown(self, timeout_millis: float = 30_000, **kwargs) -> None:
        return None

# Connect to the Platform (starts Device Provisioning if the Device is not already registered)
client = DeviceClient.start(device_id="my-device", provisioning_token="<Your Provisioning Token>", db="spotflow.db")

# Create a sender to the stream of your choice
sender = client.create_stream_sender(stream_group = "my-new-stream-group", stream = "my-new-stream")

# Initialize the custom exporter
spotflow_exporter = SpotflowMetricExporter(sender)

# Set up the periodic exporting metric reader
metric_reader = PeriodicExportingMetricReader(spotflow_exporter, export_interval_millis=1000)

# Set up the meter provider with the metric reader
meter_provider = MeterProvider(metric_readers=[metric_reader], shutdown_on_exit=False)

# Register the meter provider globally
metrics.set_meter_provider(meter_provider)

# Get a meter from the provider
meter = meter_provider.get_meter("custom-meter")

# Create a counter instrument
counter = meter.create_counter("counter")

# Record metrics in a loop
for i in range(15):
    counter.add(1)
    time.sleep(1)