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

import jakarta.annotation.PostConstruct;
import java.io.BufferedReader;
import java.io.InputStreamReader;
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: 현재 GPU 사용률 (%)
  // - ConcurrentHashMap → 멀티스레드 안전
  // =========================
  private final Map<Integer, Integer> gpuUtilMap = new ConcurrentHashMap<>();

  // =========================
  // nvidia-smi dmon 프로세스
  // - 스트리밍으로 GPU 상태를 계속 받아옴
  // =========================
  private volatile Process process;

  // =========================
  // 외부에서 GPU 사용률 조회
  // =========================
  public Map<Integer, Integer> getGpuUtilMap() {
    return gpuUtilMap;
  }

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

    if (!isNvidiaAvailable()) {
      log.warn("nvidia-smi not found. GPU monitoring disabled.");
      return;
    }

    Thread t = new Thread(this::runWithRestart, "gpu-dmon-thread");
    t.setDaemon(true); // 서버 종료 시 같이 종료
    t.start();
  }

  // =========================
  // dmon 실행 + 자동 재시작 루프
  // - dmon이 죽어도 계속 재시작
  // =========================
  private void runWithRestart() {
    boolean firstError = true;

    while (true) {
      try {
        log.info("Starting nvidia-smi dmon...");
        runDmon();
        firstError = true; // 정상 실행되면 초기화

      } catch (Exception e) {

        if (firstError) {
          log.error("nvidia-smi not available. GPU monitoring disabled.", e);
          firstError = false;
        } else {
          log.warn("dmon retry...");
        }
      }

      try {
        Thread.sleep(5000); // 5초 후에 시작
      } catch (InterruptedException ignored) {
      }
    }
  }

  // =========================
  // nvidia-smi dmon 실행 및 출력 파싱
  // =========================
  private void runDmon() throws Exception {

    // GPU utilization만 출력 (-s u)
    ProcessBuilder pb = new ProcessBuilder("nvidia-smi", "dmon", "-s", "u");

    process = pb.start();

    // dmon은 stdout으로 계속 데이터를 뿌림 (스트리밍)
    try (BufferedReader br = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
      String line;

      while ((line = br.readLine()) != null) {

        // 디버깅 로그
        // log.info("RAW: [{}]", line);

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

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

        String[] parts = line.split("\\s+");

        // GPU index 확인
        if (!parts[0].matches("\\d+")) continue;

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

        int util = 0;
        try {
          util = Integer.parseInt(parts[1]); // sm 값
        } catch (Exception e) {
          continue;
        }

        gpuUtilMap.put(index, util);
      }
    }

    // 여기 도달했다는 건 dmon 프로세스가 종료된 상태
    int exitCode = process.waitFor();
    log.warn("dmon exited. code={}", exitCode);

    // 상위 루프에서 재시작하도록 예외 발생
    throw new IllegalStateException("dmon stopped");
  }

  // =========================
  // dmon 프로세스 살아있는지 확인
  // =========================
  public boolean isAlive() {
    return process != null && process.isAlive();
  }

  // =========================
  // dmon 강제 재시작
  // - 기존 프로세스 종료 → runWithRestart에서 자동 재시작
  // =========================
  public void restart() {
    try {
      if (process != null && process.isAlive()) {
        process.destroy();
      }
    } catch (Exception ignored) {
    }
  }

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