atrade/src/main/kotlin/service/LlamaServerManager.kt

package service

import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.launch
import network.RagService
import java.io.BufferedReader
import java.io.File
import java.io.InputStreamReader
import java.util.concurrent.ConcurrentHashMap

object LlamaServerManager {
    // 포트별로 프로세스를 관리합니다.
    private val processes = ConcurrentHashMap<Int, Process>()
    private val scope = CoroutineScope(Dispatchers.IO + SupervisorJob())
    init {
        Runtime.getRuntime().addShutdownHook(Thread {
            stopAll()
        })
    }

    fun startServer(binPath: String, modelPath: String, port: Int, nGpuLayers: Int = 99) {
        // 이미 해당 포트에서 실행 중이거나 모델 경로가 비었으면 무시합니다.
        if (processes.containsKey(port) || modelPath.isBlank()) return
        val os = System.getProperty("os.name").lowercase()
        val arch = System.getProperty("os.arch").lowercase()
        val isWin = os.contains("win")
        val (nGpuLayers, threads) = when {
            os.contains("mac") && (arch.contains("arm64") || arch.contains("aarch64")) -> 99 to 8
            isWin -> 16 to 6 // NUC Core Ultra 7: GPU 레이어 40 내외, 스레드 12 권장
            else -> 0 to 4    // 인텔 맥 2017 등
        }

        val command = mutableListOf(
            binPath,
            "-m", modelPath,
            "--port", port.toString(),
            "-c", if (port == 8081) "512" else "8192",
            "-ngl", nGpuLayers.toString(),
            "-t", threads.toString(),
            "--embedding"
        )
        if (port != 8081) { // 텍스트 생성용 모델에만 적용
            command.addAll(listOf(
                "-b", "512", // Batch size (토큰 병렬 처리량 제한으로 연산 안정화)
                "--threads-batch", threads.toString(),
                "-fa","on" // Flash Attention 활성화 (메모리 절약 및 긴 컨텍스트 연산 안정성 증가)
            ))
        }
        scope.launch {
            try {
                val pb = ProcessBuilder(command)

                // 2. 윈도우 Vulkan 환경 변수 설정
                if (isWin && binPath.contains("win-x64")) {
                    val env = pb.environment()
                    // 특정 GPU 선택 (내장 GPU가 여러 개일 경우)
                    // env["GGML_VULKAN_DEVICE"] = "0"

                    // DLL 로드 경로 강제 지정 (bin 폴더 내 dll 참조)
                    val libraryPath = File(binPath).parentFile.absolutePath
                    val currentPath = System.getenv("PATH") ?: ""
                    env["PATH"] = "$libraryPath;$currentPath"

                    println("🔧 [Vulkan] 환경 변수 설정 완료: $libraryPath")
                }

                pb.redirectErrorStream(true)
                File(binPath).setExecutable(true)

                val process = pb.start()
                processes[port] = process
                println("✅ AI 서버 시작 시도 (Port: $port, Model: ${File(modelPath).name})")

                val reader = BufferedReader(InputStreamReader(process.inputStream))
                var line: String?
                while (reader.readLine().also { line = it } != null) {
                    // 로그 출력 (디버깅용)
//                     println("[Server $port] $line")
                    if (line?.contains("server is listening") == true) {
                        println("🚀 AI 서버 준비 완료 (Port: $port)")
                        if (processes.size > 1) {
                            println("[Cache] ${processes.size}")
                            RagService.active()
                        }
                    }
                }
            } catch (e: Exception) {
                println("❌ AI 서버 실행 실패 (Port: $port): ${e.message}")
                processes.remove(port)
            }

        }
    }

    fun stopAll(): Boolean {
        var allStopped = true // 모든 프로세스 종료 여부를 추적하는 플래그

        processes.forEach { (port, process) ->
            try {
                process.destroy() // 1차: 부드러운 종료 시도

                // 2차: 최대 3초 대기 후 종료되지 않으면 강제 종료
                if (!process.waitFor(3, java.util.concurrent.TimeUnit.SECONDS)) {
                    process.destroyForcibly() // 강제 사살

                    // 강제 종료 후에도 프로세스가 살아있는지 최종 확인
                    if (process.isAlive) {
                        println("❌ [Server $port] 강제 종료 명령 후에도 프로세스가 살아있습니다.")
                        allStopped = false
                    } else {
                        println("⚠️ [Server $port] 응답이 없어 강제 종료되었습니다.")
                    }
                } else {
                    println("🛑 [Server $port] 정상 종료되었습니다.")
                }
            } catch (e: Exception) {
                println("❌ [Server $port] 종료 중 오류: ${e.message}")
                allStopped = false
            }
        }

        if (allStopped) {
            processes.clear() // 모든 프로세스가 성공적으로 종료되거나 리스트에서 제거될 준비가 된 경우
        }

        return allStopped
    }
}