from plugins.base import BaseModule
import psutil
import time


class SystemMetricsCollector(BaseModule):
    def __init__(self, params):
        super().__init__(params)
        self.output_queue = None
        self.host_tag = params.get('host_tag', 'unknown')
        self.interval = params.get('interval', 1.0)

    def set_output_queue(self, queue):
        self.output_queue = queue

    def collect_metrics(self):
        return {
            'cpu_percent': psutil.cpu_percent(interval=None),
            'mem_available_gb': psutil.virtual_memory().available / (1024 ** 3),
            'disk_usage_percent': psutil.disk_usage('/').percent,
            'timestamp': int(time.time() * 1000000000)
        }

    def run(self):
        self.is_running = True
        while self.is_running:
            raw_metrics = self.collect_metrics()
            line_protocol_string = self.format_to_line_protocol(raw_metrics)

            if self.output_queue and line_protocol_string:
                # 放入 Queue，传递给 Storage 进程
                self.output_queue.put(line_protocol_string)

            time.sleep(self.interval)

    def format_to_line_protocol(self, metrics):
        tag_set = f"host={self.host_tag}"
        field_set_parts = []
        for key, value in metrics.items():
            if key == 'timestamp':
                continue
            if isinstance(value, int):
                field_set_parts.append(f"{key}={value}i")
            else:
                field_set_parts.append(f"{key}={value}")
        field_set = ",".join(field_set_parts)
        timestamp = metrics.get('timestamp')
        line = f"system_metrics,{tag_set} {field_set} {timestamp}"
        return line