smart-bow-infrared/Assets/LightGlue/Scripts/Python/PythonProcessController.cs

using System;
using System.Diagnostics;
using System.IO;
using LightGlue.Unity.Config;
using UnityEngine;
using Debug = UnityEngine.Debug;

namespace LightGlue.Unity.Python
{
    /// <summary>
    /// Unity-controlled Python child process launcher (start/stop + exit detection).
    /// This class only wraps process control; it does not change LightGlue algorithm logic.
    /// </summary>
    public sealed class PythonProcessController : MonoBehaviour
    {
        [Header("Python")]
        [Tooltip("python.exe path or just 'python' if it's on PATH.")]
        public string pythonExe = "python";

        [Tooltip("Working directory for the Python process (e.g. LightGlue_Deployment).")]
        public string workingDirectory = "";

        [Tooltip("Script path relative to workingDirectory or absolute path.")]
        public string scriptPath = "demo_lightglue_camera_position_async.py";

        [Tooltip("参考NetworkConfig.pythonScriptArgs")]
        [TextArea(3, 8)]
        public string scriptArgs = "";

        [Header("配置管理")]
        [Tooltip("是否从NetworkConfigManager自动加载Python启动参数（启动时）")]
        public bool autoLoadNetworkConfig = true;

        [Tooltip("若为 true，则从 NetworkConfig.romaPythonScriptArgs 读取并替换占位符；否则使用 NetworkConfig.pythonScriptArgs（默认旧 LightGlue_Deployment）。")]
        public bool useRomaArgsFromNetworkConfig = false;

        [Tooltip("若为 true，则使用 NetworkConfig.romaWifiPythonScriptArgs（适配 demo_roma_camera_position_async_wifi.py）；优先级高于 useRomaArgsFromNetworkConfig。")]
        public bool useRomaWifiArgsFromNetworkConfig = false;

        [Header("Lifecycle")]
        [Tooltip("是否在OnEnable时自动启动Python进程。若启用下面的「首帧后再启动」，则 OnEnable 时不启动，由 Bridge 在收到首帧图像后调用 StartPython()。")]
        public bool autoStartOnEnable = true;
        [Tooltip("为 true 时：不在一启动就拉 Python，等收到首帧图像后再启动（首次启动或应用配置重启后均生效），避免无图时白跑 TensorRT 等。")]
        public bool startOnFirstImageReceived = true;
        public bool killProcessTreeOnStop = true;

        [Header("无窗口运行（避免后台卡顿）")]
        [Tooltip("勾选后自动在启动参数中加上 --no_display，Python 不创建 OpenCV 窗口；结果仍通过 UDP 回 Unity。可避免「窗口在游戏背后时被系统节流导致卡顿」。取消勾选可保留 Python 绘制窗口便于调试。")]
        public bool addNoDisplayWhenLaunching = true;

        [Header("打包模式")]
        [Tooltip("勾选后在打包客户端中自动查找同级目录下的 LightGlue_Deployment/lightglue_runtime.exe，" +
                 "无需在 Inspector 中手动配置 pythonExe / workingDirectory / scriptPath。开发阶段建议不勾选，使用本机 Python + .py 脚本。")]
        public bool usePackagedExe = false;

        private Process _proc;
        private StreamWriter _stdinWriter;
        private volatile bool _readyToReceiveFrames;
        private float _processStartTime = -1f;

        [Tooltip("若在此秒数内未收到 Python 就绪输出，则视为就绪并开始发图（防止 stdout 缓冲导致窗口一直不弹出）")]
        public float readyTimeoutSeconds = 120f;

        public bool IsRunning => _proc != null && !_proc.HasExited;

        /// <summary>
        /// Python 是否已进入接收循环（收到脚本 stdout 中的就绪标记后才为 true，避免 TensorRT 编译期间发图导致缓冲堆积、后续解码损坏）。
        /// </summary>
        public bool IsReadyToReceiveFrames => _readyToReceiveFrames;

        /// <summary>
        /// 获取Python进程的stdin写入流（用于直接发送图片数据）
        /// </summary>
        public StreamWriter StdinWriter => _stdinWriter;

        private void Start()
        {
            // 自动加载网络配置中的Python启动参数
            if (autoLoadNetworkConfig)
            {
                LoadPythonArgsFromConfig();
            }
        }

        /// <summary>
        /// 从NetworkConfigManager加载Python启动参数
        /// </summary>
        private void LoadPythonArgsFromConfig()
        {
            try
            {
                NetworkConfig config = NetworkConfigManager.LoadConfig();
                if (config != null && config.Validate())
                {
                    // 获取处理后的参数（替换占位符）
                    string processedArgs;
                    if (useRomaWifiArgsFromNetworkConfig)
                        processedArgs = config.GetRomaWifiPythonScriptArgs();
                    else if (useRomaArgsFromNetworkConfig)
                        processedArgs = config.GetRomaPythonScriptArgs();
                    else
                        processedArgs = config.GetPythonScriptArgs();
                    if (!string.IsNullOrWhiteSpace(processedArgs))
                    {
                        scriptArgs = processedArgs;
                        Debug.Log($"[Python] 已从配置文件加载启动参数: {scriptArgs}");
                    }
                }
            }
            catch (System.Exception ex)
            {
                Debug.LogWarning($"[Python] 加载网络配置失败，使用Inspector中的默认值: {ex.Message}");
            }
        }

        private void OnEnable()
        {
            if (autoStartOnEnable && !startOnFirstImageReceived)
                StartPython();
        }

        private void OnDisable()
        {
            StopPython();
        }

        private void Update()
        {
            if (IsRunning && !_readyToReceiveFrames && _processStartTime >= 0f && readyTimeoutSeconds > 0f
                && (UnityEngine.Time.time - _processStartTime) >= readyTimeoutSeconds)
            {
                _readyToReceiveFrames = true;
                Debug.Log("[Python] 未检测到就绪输出，超时后视为就绪并开始发送。");
            }
        }

        public void StartPython()
        {
            if (IsRunning) return;

            // 确保在真正启动进程前，先加载最新的配置参数
            if (autoLoadNetworkConfig)
            {
                LoadPythonArgsFromConfig();
            }

            string wd = workingDirectory;
            string packagedExePath = null;

            if (!usePackagedExe)
            {
                // 开发模式：使用本机 Python + 脚本.py
                if (string.IsNullOrWhiteSpace(wd))
                    wd = Directory.GetCurrentDirectory();
                wd = Path.GetFullPath(wd);
            }
            else
            {
                // 打包模式：自动定位 LightGlue_Deployment 目录和 lightglue_runtime.exe
                string baseDir;
#if UNITY_STANDALONE
                // 打包后的客户端：Application.dataPath 指向 xxx_Data，父目录为 exe 所在目录
                baseDir = Path.GetDirectoryName(Application.dataPath);
#else
                // 编辑器或其他平台：使用当前工作目录，方便在本机模拟客户端目录结构
                baseDir = Directory.GetCurrentDirectory();
#endif
                string deploymentDir = Path.Combine(baseDir, "LightGlue_Deployment");
                packagedExePath = Path.Combine(deploymentDir, "lightglue_runtime.exe");

                wd = deploymentDir;

                if (!File.Exists(packagedExePath))
                {
                    Debug.LogError($"[Python] 打包模式启动失败：未找到打包好的运行文件: {packagedExePath}");
                    return;
                }
            }

            string script = scriptPath;
            if (!usePackagedExe)
            {
                // 开发模式：仍然使用 python + 脚本.py 方式启动
                if (!Path.IsPathRooted(script))
                    script = Path.GetFullPath(Path.Combine(wd, script));

                if (!File.Exists(script))
                {
                    Debug.LogError($"[Python] Script not found: {script}");
                    return;
                }
            }

            var psi = new ProcessStartInfo
            {
                // 打包模式：直接执行打包好的 exe；开发模式：执行 python.exe
                FileName = usePackagedExe ? packagedExePath : pythonExe,
                WorkingDirectory = wd,
                UseShellExecute = false,
                CreateNoWindow = true,
                RedirectStandardOutput = true,
                RedirectStandardError = true,
                RedirectStandardInput = true,  // 启用stdin重定向
            };
            string args = (scriptArgs ?? string.Empty).Trim();
            if (addNoDisplayWhenLaunching && !args.Contains("--no_display"))
                args = args + " --no_display";
            psi.Arguments = usePackagedExe ? args : $"{Quote(script)} {args}";

            _proc = new Process { StartInfo = psi, EnableRaisingEvents = true };
            _readyToReceiveFrames = false;
            _proc.OutputDataReceived += (_, e) =>
            {
                if (string.IsNullOrEmpty(e.Data)) return;
                if (!_readyToReceiveFrames && (e.Data.Contains("UDP JPEG mode: receiver started") || e.Data.Contains("First frame received and processed")))
                {
                    _readyToReceiveFrames = true;
                    Debug.Log("[Python] 已就绪接收图像，开始向 Python 发送帧。");
                }
                Debug.Log($"[Python] {e.Data}");
            };
            _proc.ErrorDataReceived += (_, e) =>
            {
                if (string.IsNullOrEmpty(e.Data)) return;
                // Some Python libraries output warnings to stderr that are not fatal errors
                string msg = e.Data.ToLower();
                if (msg.Contains("unable to import") || msg.Contains("please install") || 
                    msg.Contains("warning") || msg.Contains("deprecated"))
                {
                    Debug.LogWarning($"[Python] {e.Data}");
                }
                else
                {
                    Debug.LogError($"[Python] {e.Data}");
                }
            };
            _proc.Exited += (_, __) =>
            {
                int code = _proc?.ExitCode ?? -1;
                Debug.LogWarning($"[Python] Process exited. code={code}");
            };

            try
            {
                if (!_proc.Start())
                {
                    Debug.LogError("[Python] Failed to start process.");
                    _proc = null;
                    return;
                }
                
                // 获取stdin写入流（用于发送图片数据）
                _stdinWriter = _proc.StandardInput;
                _stdinWriter.AutoFlush = true;  // 自动刷新，减少延迟
                
                _processStartTime = UnityEngine.Time.time;
                _proc.BeginOutputReadLine();
                _proc.BeginErrorReadLine();
                Debug.Log($"[Python] Started successfully! pid={_proc.Id}\n" +
                         $"  WorkingDir: {wd}\n" +
                         $"  Command: {psi.FileName} {psi.Arguments}\n" +
                         $"  Mode: {(usePackagedExe ? "PackagedExe" : "DevPython")}\n" +
                         $"  Stdin: Enabled (for direct image transfer)");
            }
            catch (Exception ex)
            {
                Debug.LogError($"[Python] Start failed: {ex}");
                _proc = null;
                _stdinWriter = null;
            }
        }

        public void StopPython()
        {
            if (_proc == null) return;

            try
            {
                if (!_proc.HasExited)
                {
                    if (killProcessTreeOnStop)
                        KillProcessTree(_proc);
                    else
                        _proc.Kill();
                }
            }
            catch (Exception ex)
            {
                Debug.LogWarning($"[Python] Stop error: {ex.Message}");
            }
            finally
            {
                _readyToReceiveFrames = false;
                _processStartTime = -1f;
                try 
                { 
                    _stdinWriter?.Close(); 
                    _stdinWriter = null;
                } 
                catch { /* ignore */ }
                try { _proc.Dispose(); } catch { /* ignore */ }
                _proc = null;
            }
        }

        /// <summary>
        /// Attempts to kill the entire process tree for the given process.
        /// On Windows this uses 'taskkill /T /F', on other platforms it falls back to Kill().
        /// </summary>
        private static void KillProcessTree(Process proc)
        {
            if (proc == null) return;

            try
            {
#if UNITY_STANDALONE_WIN || UNITY_EDITOR_WIN
                // Use Windows taskkill to terminate the whole tree
                try
                {
                    var startInfo = new ProcessStartInfo
                    {
                        FileName = "taskkill",
                        Arguments = $"/PID {proc.Id} /T /F",
                        CreateNoWindow = true,
                        UseShellExecute = false,
                        RedirectStandardOutput = false,
                        RedirectStandardError = false
                    };
                    using (var killer = Process.Start(startInfo))
                    {
                        // Best-effort: wait a short time, but don't block indefinitely
                        killer?.WaitForExit(3000);
                    }
                }
                catch
                {
                    // If taskkill fails, fall back to normal Kill
                    if (!proc.HasExited)
                        proc.Kill();
                }
#else
                // Non-Windows: no simple built-in tree-kill; just kill the main process
                if (!proc.HasExited)
                    proc.Kill();
#endif
            }
            catch
            {
                // Swallow any errors; caller already wraps StopPython() in a try/catch.
                if (!proc.HasExited)
                    try { proc.Kill(); } catch { /* ignore */ }
            }
        }

        private static string Quote(string s)
        {
            if (string.IsNullOrEmpty(s)) return "\"\"";
            if (s.Contains("\"")) s = s.Replace("\"", "\\\"");
            if (s.Contains(" ") || s.Contains("\t")) return $"\"{s}\"";
            return s;
        }
    }
}