from influxdb_client import InfluxDBClient, Point
from influxdb_client.client.write_api import SYNCHRONOUS
from loguru import logger
import os
from datetime import datetime
from typing import List, Dict

class InfluxDBManager:
    def __init__(self):
        self.url = os.getenv('INFLUXDB_URL', 'http://localhost:8086')
        self.token = os.getenv('INFLUXDB_TOKEN')
        self.org = os.getenv('INFLUXDB_ORG')
        self.bucket = os.getenv('INFLUXDB_BUCKET')
        
        if not all([self.url, self.token, self.org, self.bucket]):
            raise ValueError("InfluxDB 설정이 완료되지 않았습니다.")
            
        self.client = InfluxDBClient(
            url=self.url,
            token=self.token,
            org=self.org
        )
        self.write_api = self.client.write_api(write_options=SYNCHRONOUS)
        self.query_api = self.client.query_api()
        logger.info("InfluxDB 클라이언트 초기화 완료")

    def write_gpu_metrics(self, gpu_data: dict):
        """GPU 메트릭을 InfluxDB에 저장합니다."""
        try:
            point = Point("gpu_metrics") \
                .tag("gpu_id", str(gpu_data['id'])) \
                .field("temperature", gpu_data['temperature']) \
                .field("utilization", gpu_data['utilization']) \
                .field("memory_used", gpu_data['memory_used']) \
                .field("memory_total", gpu_data['memory_total']) \
                .field("power_draw", gpu_data['power_draw']) \
                .field("power_limit", gpu_data['power_limit']) \
                .time(datetime.utcnow())
                
            self.write_api.write(bucket=self.bucket, record=point)
            
        except Exception as e:
            logger.error(f"GPU 메트릭 저장 실패: {str(e)}")

    def write_cooling_metrics(self, cooling_data: dict):
        """냉각 시스템 메트릭을 InfluxDB에 저장합니다."""
        try:
            for sensor_name, temp in cooling_data.items():
                point = Point("cooling_metrics") \
                    .tag("sensor", sensor_name) \
                    .field("temperature", temp) \
                    .time(datetime.utcnow())
                    
                self.write_api.write(bucket=self.bucket, record=point)
                
        except Exception as e:
            logger.error(f"냉각 메트릭 저장 실패: {str(e)}")

    def get_metrics_history(self, hours: int = 24) -> List[Dict]:
        """지정된 시간 범위의 메트릭 히스토리를 조회합니다."""
        try:
            query = f'''
                from(bucket: "{self.bucket}")
                    |> range(start: -{hours}h)
                    |> filter(fn: (r) => r["_measurement"] == "gpu_metrics" or r["_measurement"] == "cooling_metrics")
            '''
            
            result = self.query_api.query(query)
            
            metrics = []
            for table in result:
                for record in table.records:
                    metrics.append({
                        '_measurement': record.get_measurement(),
                        '_time': record.get_time().isoformat(),
                        '_value': record.get_value(),
                        'tags': record.values
                    })
            
            return metrics
        except Exception as e:
            logger.error(f"메트릭 히스토리 조회 실패: {str(e)}")
            return []

    def close(self):
        """InfluxDB 연결을 종료합니다."""
        try:
            self.write_api.close()
            self.client.close()
            logger.info("InfluxDB 연결 종료")
        except Exception as e:
            logger.error(f"InfluxDB 연결 종료 실패: {str(e)}")

    def __del__(self):
        self.close() 