package com.kamco.cd.training.common.service;

import jakarta.annotation.PostConstruct;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayDeque;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import lombok.extern.log4j.Log4j2;
import org.springframework.stereotype.Component;

@Component
@Log4j2
public class GpuDmonReader {

  // =========================
  // GPU 사용률 히스토리 저장소
  // key: GPU index (0,1,2...)
  // value: 최근 WINDOW_SIZE 개의 sm 사용률 샘플 (1초 간격)
  // ConcurrentHashMap → 멀티스레드 안전 (deque 접근은 synchronized)
  // =========================
  private static final int WINDOW_SIZE = 10; // 10초 평균

  private final Map<Integer, ArrayDeque<Integer>> historyMap = new ConcurrentHashMap<>();

  // =========================
  // 외부 조회용
  // GPU 개수만큼 10초 평균값 반환
  // SystemMonitorService에서 호출
  // =========================
  public Map<Integer, Integer> getGpuUtilMap() {
    Map<Integer, Integer> result = new HashMap<>();
    historyMap.forEach(
        (gpu, deque) -> {
          synchronized (deque) {
            int avg =
                deque.isEmpty()
                    ? 0
                    : (int)
                        Math.round(deque.stream().mapToInt(Integer::intValue).average().orElse(0));
            result.put(gpu, avg);
          }
        });
    return result;
  }

  // =========================
  // Bean 초기화 시 실행
  // - 별도 스레드에서 GPU 모니터링 시작
  // - 메인 스레드 block 방지
  // =========================
  @PostConstruct
  public void start() {

    // nvidia-smi 없는 환경이면 GPU 모니터링 비활성화
    if (!isNvidiaAvailable()) {
      log.warn("nvidia-smi not found. GPU monitoring disabled.");
      return;
    }

    // 데몬 스레드로 실행 (서버 종료 시 자동 종료)
    Thread t = new Thread(this::runLoop, "gpu-dmon-thread");
    t.setDaemon(true);
    t.start();
  }

  // =========================
  // 무한 루프
  // - dmon 실행
  // - 죽으면 자동 재시작
  // =========================
  private void runLoop() {

    while (true) {
      try {
        runDmon(); // GPU 사용률 수집 시작
      } catch (Exception e) {
        // dmon 프로세스 종료되면 여기로 들어옴
        log.warn("dmon restart: {}", e.getMessage());
      }

      // 5초 대기 후 재시작
      sleep(5000);
    }
  }

  // =========================
  // nvidia-smi dmon 실행
  // - GPU 사용률 스트리밍으로 계속 수신
  // =========================
  private void runDmon() throws Exception {

    // -s u → GPU utilization만 출력
    ProcessBuilder pb = new ProcessBuilder("nvidia-smi", "dmon", "-s", "u");
    // stderr를 stdout에 합쳐서 소비 — 소비하지 않으면 stderr 버퍼가 가득 차
    // nvidia-smi 프로세스 자체가 block되어 stdout도 멈춤
    pb.redirectErrorStream(true);

    Process process = pb.start();
    // 프로세스 실행 후 stdout 읽기
    try (BufferedReader br = new BufferedReader(new InputStreamReader(process.getInputStream()))) {

      String line;

      // dmon은 계속 출력됨 (스트리밍)
      while ((line = br.readLine()) != null) {

        // 헤더 제거 (#로 시작)
        if (line.startsWith("#")) continue;

        line = line.trim();
        if (line.isEmpty()) continue;

        // 공백 기준 분리
        String[] parts = line.split("\\s+");

        // 첫 번째 값이 GPU index인지 확인
        if (parts.length < 2 || !parts[0].matches("\\d+")) continue;

        int index = Integer.parseInt(parts[0]);

        try {
          // 두 번째 값이 GPU 사용률 (sm), "-"이면 N/A (GPU 미활성)
          int util = Integer.parseInt(parts[1]);
          pushSample(index, util);
          log.debug("[GpuDmon] gpu={} sm={}%", index, util);

        } catch (NumberFormatException e) {
          // "-" 등 숫자가 아닌 값 → GPU 비활성 상태이므로 0으로 기록
          pushSample(index, 0);
          log.debug("[GpuDmon] gpu={} sm=N/A (raw={}), stored as 0", index, parts[1]);
        }
      }
    } finally {
      process.destroy();
    }

    // 여기까지 왔다는 건 dmon 프로세스 종료됨
    // → runLoop에서 재시작하도록 예외 발생
    throw new IllegalStateException("dmon stopped");
  }

  // =========================
  // 샘플 추가 — 윈도우 초과 시 가장 오래된 값 제거
  // =========================
  private void pushSample(int gpuIndex, int util) {
    ArrayDeque<Integer> deque =
        historyMap.computeIfAbsent(gpuIndex, k -> new ArrayDeque<>(WINDOW_SIZE));
    synchronized (deque) {
      deque.addLast(util);
      if (deque.size() > WINDOW_SIZE) {
        deque.pollFirst();
      }
    }
  }

  // =========================
  // nvidia-smi 존재 여부 확인
  // =========================
  private boolean isNvidiaAvailable() {
    try {
      Process p = new ProcessBuilder("which", "nvidia-smi").start();
      return p.waitFor() == 0;
    } catch (Exception e) {
      return false;
    }
  }

  // =========================
  // sleep 유틸
  // =========================
  private void sleep(long ms) {
    try {
      Thread.sleep(ms);
    } catch (InterruptedException ignored) {
    }
  }
}
