package connect

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"FireLeave_tool/logger"
)

var (
	InferenceProcesses     sync.Map // key: process UUID, value: *exec.Cmd
	inferenceLastHeartbeat sync.Map // key: process UUID, value: time.Time
)

// UpdateInferenceHeartbeat records the latest heartbeat time for one inference process.
func UpdateInferenceHeartbeat(processUUID string) {
	inferenceLastHeartbeat.Store(processUUID, time.Now())
	if strings.Contains(processUUID, "_fire_check") {
		logger.Logger.Printf("推理心跳收到 | 时间: %s | UUID: %s | 类型: fire_check", time.Now().Format("2006/01/02 15:04:05.000"), processUUID)
	} else {
		logger.Logger.Printf("推理心跳收到 | 时间: %s | UUID: %s | 类型: fire_leave", time.Now().Format("2006/01/02 15:04:05.000"), processUUID)
	}
}

// ShouldRestartInference reports whether one inference process should be restarted.
func ShouldRestartInference(processUUID string) bool {
	val, ok := inferenceLastHeartbeat.Load(processUUID)
	if !ok {
		return true
	}
	last := val.(time.Time)
	return time.Since(last) > 85*time.Second
}

// StopInferenceProcess stops one inference process.
func StopInferenceProcess(processUUID string) {
	if raw, ok := InferenceProcesses.Load(processUUID); ok {
		if cmd, ok := raw.(*exec.Cmd); ok && cmd.Process != nil {
			logger.Logger.Printf("正在停止推理进程 (ProcessUUID=%s, PID=%d)", processUUID, cmd.Process.Pid)
			_ = cmd.Process.Kill()
			_ = cmd.Wait()
		}
		InferenceProcesses.Delete(processUUID)
	}
}

// StartInferenceProcess keeps the old "restart both" behavior for compatibility.
func StartInferenceProcess(deviceData *DeviceData) {
	if err := RestartInferenceProcess(deviceData, deviceData.DeviceUUID); err != nil {
		logger.Logger.Printf("fire_leave推理程序重启失败 (DeviceUUID=%s): %v", deviceData.DeviceUUID, err)
	}
	if err := RestartInferenceProcess(deviceData, fmt.Sprintf("%s_fire_check", deviceData.DeviceUUID)); err != nil {
		logger.Logger.Printf("fire_check推理程序重启失败 (DeviceUUID=%s): %v", deviceData.DeviceUUID, err)
	}
}

// RestartInferenceProcess restarts only the requested inference process UUID.
func RestartInferenceProcess(deviceData *DeviceData, processUUID string) error {
	if deviceData == nil {
		return fmt.Errorf("device data is nil")
	}

	port, err := GlobalPortManager.GetPort(processUUID)
	if err != nil {
		return fmt.Errorf("get port failed: %w", err)
	}

	modelName := "model"
	detectAreaStr := ProcessDetectAreaForInference(deviceData.FireLeaveDetectArea)
	args := []string{
		"-s", deviceData.CameraRTSP,
		"-m", modelName,
		"-c", processUUID,
		"-p", fmt.Sprintf("%d", port),
		"-w", detectAreaStr,
		"-r", fmt.Sprintf("%d", deviceData.Confidence),
	}

	if strings.HasSuffix(processUUID, "_fire_check") {
		modelName = "fire_check"
		detectAreaStr = ProcessDetectAreaForInference(deviceData.FireCheckDetectArea)
		confidence := GetFireCheckConfidence(*deviceData)
		args = []string{
			"-s", deviceData.CameraRTSP,
			"-m", modelName,
			"-c", processUUID,
			"-p", fmt.Sprintf("%d", port),
			"-w", detectAreaStr,
			"-r", fmt.Sprintf("%d", confidence),
		}
	}

	StopInferenceProcess(processUUID)
	cmd := exec.Command("./yolov5", args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	logger.Logger.Printf("正在重启推理程序 (ProcessUUID=%s, Port=%d): %v", processUUID, port, cmd.Args)
	if err := cmd.Start(); err != nil {
		return err
	}

	InferenceProcesses.Store(processUUID, cmd)
	UpdateInferenceHeartbeat(processUUID)
	logger.Logger.Printf("推理程序已成功重启，PID=%d (ProcessUUID=%s, Port=%d)", cmd.Process.Pid, processUUID, port)
	return nil
}