import os

from .common import AbstractMetricCollector
from prometheus_client import Gauge
import gc
from loguru import logger


class PythonGcMetricsCollector(AbstractMetricCollector):
    def __init__(self, interval_sec: int = 30):
        super().__init__(interval_sec=interval_sec)
        self.pid = os.getpid()
        # GC 总次数
        self.gc_collections_total = Gauge(
            "app_python_gc_collections_total",
            "Total number of GC collections",
            ["generation", "pid"],
        )
        # 累计回收的对象总数
        self.gc_collected_total = Gauge(
            "app_python_gc_collected_total",
            "Total objects collected",
            ["generation", "pid"],
        )
        # 累计无法回收的对象总数
        self.gc_uncollectable_total = Gauge(
            "app_python_gc_uncollectable_total",
            "Total uncollectable objects",
            ["generation", "pid"],
        )

    async def collect(self):
        try:
            # 获取当前统计信息
            stats = gc.get_stats()

            # 更新指标
            for gen, stat in enumerate(stats):
                self.gc_collections_total.labels(generation=gen, pid=self.pid).set(
                    stat.get("collections", 0)
                )
                self.gc_collected_total.labels(generation=gen, pid=self.pid).set(
                    stat.get("collected", 0)
                )
                self.gc_uncollectable_total.labels(generation=gen, pid=self.pid).set(
                    stat.get("uncollectable", 0)
                )

        except Exception as e:
            logger.error(f"Error collecting GC metrics: {e}")


class PythonGcObjectMetricsCollector(AbstractMetricCollector):
    def __init__(self, interval_sec: int = 30):
        super().__init__(interval_sec=interval_sec)
        self.pid = os.getpid()
        # 当前时刻回收的对象数
        self.gc_objects_count = Gauge(
            "app_python_gc_objects_count",
            "Current number of objects",
            ["generation", "pid"],
        )

    async def collect(self):
        try:
            counts = gc.get_count()
            for gen, count in enumerate(counts):
                self.gc_objects_count.labels(generation=gen, pid=self.pid).set(count)

        except Exception as e:
            logger.error(f"Error collecting GC Object metrics: {e}")
