import asyncio
import logging
import os
import threading
from abc import ABC
from typing import List, Type

import psutil
from prometheus_client import Gauge

from .common import AbstractMetricCollector

logger = logging.getLogger(__name__)


class AbstractSystemMetricCollector(AbstractMetricCollector, ABC):
    """系统指标收集器抽象基类"""

    def __init__(self, process: psutil.Process, pid: int, **kwargs):
        super().__init__()
        self.process = process
        self.pid = pid


class CPUMetricCollector(AbstractSystemMetricCollector):
    """CPU使用率指标收集器"""

    def __init__(self, process: psutil.Process, pid: int):
        super().__init__(process, pid)
        self.gauge = Gauge(
            "cpu_usage_percent",
            "CPU usage in percent",
            ["pid"],
            multiprocess_mode="all",
        )

    async def collect(self):
        cpu_usage = self.process.cpu_percent()
        self.gauge.labels(pid=self.pid).set(cpu_usage)


class MemoryMetricCollector(AbstractSystemMetricCollector):
    """内存使用量指标收集器"""

    def __init__(self, process: psutil.Process, pid: int):
        super().__init__(process, pid)
        self.gauge_rss = Gauge(
            "resident_memory_usage_mb",
            "Resident Memory usage in MB",
            ["pid"],
            multiprocess_mode="all",
        )
        self.gauge_vms = Gauge(
            "virtual_memory_usage_mb",
            "Virtual Memory usage in MB",
            ["pid"],
            multiprocess_mode="all",
        )

    async def collect(self):
        memory_info = self.process.memory_info()
        rss = memory_info.rss / (1024 * 1024)  # resident set size resident set size
        vms = memory_info.vms / (1024 * 1024)  # virtual memory size
        self.gauge_rss.labels(pid=self.pid).set(rss)
        self.gauge_vms.labels(pid=self.pid).set(vms)


class ThreadCountMetricCollector(AbstractSystemMetricCollector):
    """线程数指标收集器"""

    def __init__(self, process: psutil.Process, pid: int):
        super().__init__(process, pid)
        self.gauge_active_thread_count = Gauge(
            "active_thread_count",
            "Number of active threads",
            ["pid"],
            multiprocess_mode="all",
        )
        self.gauge_pid_thread_count = Gauge(
            "thread_count", "Number of Pid threads", ["pid"], multiprocess_mode="all"
        )

    async def collect(self):
        pid_thread_count = self.process.num_threads()
        self.gauge_pid_thread_count.labels(pid=self.pid).set(pid_thread_count)
        self.gauge_active_thread_count.labels(pid=self.pid).set(
            threading.active_count()
        )


class FileMetricCollector(AbstractSystemMetricCollector):
    """文件描述符数指标收集器"""

    def __init__(self, process: psutil.Process, pid: int):
        super().__init__(process, pid)
        self.gauge = Gauge(
            "open_fds",
            "Number of open files",
            ["pid"],
            multiprocess_mode="all",
        )

    async def collect(self):
        file_descriptor_count = len(self.process.open_files())
        self.gauge.labels(pid=self.pid).set(file_descriptor_count)


class SystemMetricsMonitor(AbstractMetricCollector):
    """系统指标监控器 - 采用组装模式"""

    def __init__(self, interval_sec: int = 30):
        super().__init__(interval_sec=interval_sec)
        self.pid = os.getpid()
        self.process = psutil.Process(self.pid)
        self.collectors: List[AbstractSystemMetricCollector] = []

    def add(
        self, collector_classes: Type[AbstractSystemMetricCollector], **kwargs
    ) -> "SystemMetricsMonitor":
        self.collectors.append(collector_classes(self.process, self.pid, **kwargs))
        return self

    def cpu(self):
        return self.add(CPUMetricCollector)

    def memory(self):
        return self.add(MemoryMetricCollector)

    def thread(self):
        return self.add(ThreadCountMetricCollector)

    def file(self):
        return self.add(FileMetricCollector)

    def register_default_collectors(self) -> "SystemMetricsMonitor":
        """
        组装默认指标收集器
        """
        self.add(CPUMetricCollector)
        self.add(MemoryMetricCollector)
        self.add(ThreadCountMetricCollector)
        self.add(FileMetricCollector)
        return self

    async def collect(self):
        for collector in self.collectors:
            await collector.collect()

    async def run_monitor(self, interval_sec: int = 10):
        """运行监控任务"""

        self.set_interval(interval_sec)

        async def work():
            while True:
                # 执行所有收集器的收集任务
                for collector in self.collectors:
                    await collector.collect()

                await asyncio.sleep(interval_sec)

        if len(self.collectors) == 0:
            self.register_default_collectors()

        logger.info(
            f"Start system metrics monitor, interval: {interval_sec}s, collector num = {len(self.collectors)}"
        )

        asyncio.create_task(work())
