LightGlueUnity/Assets/LightGlue/Scripts/Sdk/Core/LightGlueClient.cs

using System;
using System.IO;
using System.Net;
using System.Net.Sockets;
using LightGlue.Unity.Config;
using LightGlue.Unity.Networking;
using LightGlue.Unity.Python;
using LightGlue.Unity.Sdk.Core.Networking;
using UnityEngine;

namespace LightGlue.Unity.Sdk.Core
{
    /// <summary>
    /// SDK 原型版：对外统一的 LightGlue 通信客户端。
    /// - 组合 JPEG 接收、Python 结果接收、硬件控制 0x40 发送、Unity->Python 图像发送。
    /// - 非 MonoBehaviour，由外部（如 BridgeBehaviour）驱动生命周期和 Tick。
    /// </summary>
    public sealed class LightGlueClient : IDisposable
    {
        public enum TransferMode
        {
            Stdin,
            Udp
        }

        public sealed class Options
        {
            // Hardware -> Unity image stream
            public string HardwareBindIp = "0.0.0.0";
            public int HardwarePort = 12346;
            public float HardwareTimeoutSeconds = 2.0f;
            public int MaxQueuedFrames = 2;

            // Unity -> Hardware control (0x40)
            public string HardwareControlIp = "192.168.0.106";
            public int HardwareControlPort = 8888;

            // Unity -> Python image
            public TransferMode Mode = TransferMode.Stdin;
            public string PythonIp = "127.0.0.1";
            public int PythonPort = 12347;
            public PythonProcessController PythonController;

            // Python -> Unity result
            public bool EnableResultReceiver = true;
            public string PythonResultBindIp = "127.0.0.1";
            public int PythonResultPort = 12348;
            public int MaxResultQueueSize = 10;
        }

        private readonly Options _options;

        private CoreUDPJpegReceiver _jpegReceiver;
        private CoreUDPResultReceiver _resultReceiver;

        private UdpClient _pythonUdpSender;
        private IPEndPoint _pythonEndpoint;

        private UdpClient _hwControlClient;
        private IPEndPoint _hwControlEndpoint;

        private PythonProcessController _pythonController;
        private StreamWriter _stdinWriter;

        private ImageTransmissionConfig _transmissionConfig;
        private bool _configEnabled = true;

        private byte[] _latestJpeg;
        private readonly object _jpegLock = new object();
        private float _lastSendTimeMs;

        public bool IsRunning { get; private set; }

        public LightGlueClient(Options options)
        {
            _options = options ?? throw new ArgumentNullException(nameof(options));
        }

        /// <summary>
        /// 启动所有内部模块（JPEG 接收、结果接收、硬件控制 UDP、Python 图像发送端）。
        /// </summary>
        public void Start()
        {
            if (IsRunning) return;

            // JPEG receiver
            _jpegReceiver = new CoreUDPJpegReceiver(
                _options.HardwareBindIp,
                _options.HardwarePort,
                _options.HardwareTimeoutSeconds,
                _options.MaxQueuedFrames);
            try
            {
                _jpegReceiver.Start();
            }
            catch (SocketException ex)
            {
                Debug.LogError($"[SDK][Client] Failed to bind JPEG receiver on {_options.HardwareBindIp}:{_options.HardwarePort}: {ex.Message}");
                SafeStop(ref _jpegReceiver);
                throw;
            }

            // Unity -> Python sender
            if (_options.Mode == TransferMode.Stdin)
            {
                _pythonController = _options.PythonController;
                TryConnectStdin();
            }
            else
            {
                var ip = IPAddress.Parse(string.IsNullOrWhiteSpace(_options.PythonIp) ? "127.0.0.1" : _options.PythonIp);
                _pythonEndpoint = new IPEndPoint(ip, _options.PythonPort);
                _pythonUdpSender = new UdpClient();
                Debug.Log($"[SDK][Client] Unity -> Python UDP sender to {_pythonEndpoint.Address}:{_pythonEndpoint.Port}");
            }

            // Hardware control (0x40)
            try
            {
                string ctrlIp = string.IsNullOrWhiteSpace(_options.HardwareControlIp)
                    ? "192.168.0.106"
                    : _options.HardwareControlIp;
                var ip = IPAddress.Parse(ctrlIp);
                _hwControlEndpoint = new IPEndPoint(ip, _options.HardwareControlPort);
                _hwControlClient = new UdpClient();
                Debug.Log($"[SDK][Client] Hardware control UDP -> {_hwControlEndpoint.Address}:{_hwControlEndpoint.Port}");
            }
            catch (Exception ex)
            {
                Debug.LogError($"[SDK][Client] Hardware control UDP init failed: {ex.Message}");
                _hwControlClient = null;
            }

            // Result receiver
            if (_options.EnableResultReceiver)
            {
                try
                {
                    _resultReceiver = new CoreUDPResultReceiver(
                        _options.PythonResultBindIp,
                        _options.PythonResultPort,
                        _options.MaxResultQueueSize);
                    _resultReceiver.Start();
                    Debug.Log($"[SDK][Client] Result receiver on {_options.PythonResultBindIp}:{_options.PythonResultPort}");
                }
                catch (Exception ex)
                {
                    Debug.LogError($"[SDK][Client] Result receiver init failed: {ex.Message}");
                    SafeStop(ref _resultReceiver);
                    _resultReceiver = null;
                }
            }

            IsRunning = true;
        }

        public void Stop()
        {
            if (!IsRunning) return;

            SafeStop(ref _jpegReceiver);
            SafeStop(ref _resultReceiver);

            try { _hwControlClient?.Close(); } catch { /* ignore */ }
            _hwControlClient = null;

            try { _pythonUdpSender?.Close(); } catch { /* ignore */ }
            _pythonUdpSender = null;

            _stdinWriter = null;
            _pythonController = null;

            IsRunning = false;
        }

        private static void SafeStop<T>(ref T disposable) where T : class, IDisposable
        {
            if (disposable == null) return;
            try { disposable.Dispose(); } catch { /* ignore */ }
            disposable = null;
        }

        private void TryConnectStdin()
        {
            if (_pythonController == null)
            {
                Debug.LogWarning("[SDK][Client] PythonProcessController is null, cannot use stdin mode.");
                return;
            }

            if (_pythonController.IsRunning)
            {
                _stdinWriter = _pythonController.StdinWriter;
                if (_stdinWriter != null)
                {
                    Debug.Log("[SDK][Client] Connected to Python stdin.");
                }
                else
                {
                    Debug.LogWarning("[SDK][Client] Python stdin writer is null.");
                }
            }
            else
            {
                Debug.LogWarning("[SDK][Client] Python process not running, stdin mode will be retried later.");
            }
        }

        /// <summary>
        /// 设置图像传输配置（仅影响 Unity->Python 图像流，不影响硬件 0x40 配置）。
        /// </summary>
        public void SetTransmissionConfig(ImageTransmissionConfig config)
        {
            _transmissionConfig = config;
            _configEnabled = config != null && config.enableImageTransmission;
        }

        /// <summary>
        /// 下发硬件图像参数（0x40 帧），通常在 UI 点击“应用配置”时调用。
        /// </summary>
        public void ApplyHardwareConfig(ImageTransmissionConfig config)
        {
            if (_hwControlClient == null || _hwControlEndpoint == null || config == null)
            {
                Debug.LogWarning("[SDK][Client] Hardware control not ready or config null, cannot send 0x40 frame.");
                return;
            }

            try
            {
                byte[] frame = CoreHardwareControlProtocol.BuildParameterSettingFrame(config);
                if (frame != null && frame.Length > 0)
                {
                    _hwControlClient.Send(frame, frame.Length, _hwControlEndpoint);
                    Debug.Log($"[SDK][Client] Sent 0x40 frame -> {_hwControlEndpoint.Address}:{_hwControlEndpoint.Port}");
                }
            }
            catch (SocketException ex)
            {
                Debug.LogWarning($"[SDK][Client] Hardware control UDP send error: {ex.SocketErrorCode} {ex.Message}");
            }
            catch (Exception ex)
            {
                Debug.LogWarning($"[SDK][Client] Hardware control send failed: {ex.Message}");
            }
        }

        /// <summary>
        /// 每帧在 Unity 的 Update 中调用，用于驱动 JPEG 队列消费并发送到 Python。
        /// </summary>
        public void Tick()
        {
            if (!IsRunning || _jpegReceiver == null) return;

            // 确保发送端可用
            if (_options.Mode == TransferMode.Stdin)
            {
                if (_stdinWriter == null)
                {
                    if (_pythonController != null && _pythonController.IsRunning)
                    {
                        _stdinWriter = _pythonController.StdinWriter;
                    }
                    if (_stdinWriter == null) return;
                }
            }
            else
            {
                if (_pythonUdpSender == null || _pythonEndpoint == null) return;
            }

            bool shouldTransmit = _configEnabled;
            if (_transmissionConfig != null)
            {
                shouldTransmit = _transmissionConfig.enableImageTransmission;
            }

            if (!shouldTransmit)
            {
                while (_jpegReceiver.TryDequeueJpeg(out _)) { }
                return;
            }

            float nowMs = Time.time * 1000f;
            if (_transmissionConfig != null)
            {
                float intervalMs = _transmissionConfig.reportIntervalMs;
                if (intervalMs > 0 && (nowMs - _lastSendTimeMs) < intervalMs)
                {
                    while (_jpegReceiver.TryDequeueJpeg(out _)) { }
                    return;
                }
            }

            byte[] latest = null;
            while (_jpegReceiver.TryDequeueJpeg(out var jpeg))
            {
                latest = jpeg;
            }

            if (latest == null) return;

            lock (_jpegLock)
            {
                _latestJpeg = latest;
            }

            byte[] processed = ProcessImage(latest);

            try
            {
                if (_options.Mode == TransferMode.Stdin)
                {
                    if (_stdinWriter != null)
                    {
                        _stdinWriter.BaseStream.Write(processed, 0, processed.Length);
                        _stdinWriter.BaseStream.Flush();
                        _lastSendTimeMs = nowMs;
                    }
                }
                else
                {
                    _pythonUdpSender.Send(processed, processed.Length, _pythonEndpoint);
                    _lastSendTimeMs = nowMs;
                }
            }
            catch (SocketException ex)
            {
                Debug.LogWarning($"[SDK][Client] UDP send error: {ex.SocketErrorCode} {ex.Message}");
            }
            catch (IOException ex)
            {
                Debug.LogWarning($"[SDK][Client] Stdin write error: {ex.Message}");
                _stdinWriter = null;
            }
        }

        private byte[] ProcessImage(byte[] jpegBytes)
        {
            if (_transmissionConfig == null)
                return jpegBytes;

            // 目前保持与现有 Bridge 一致：不在 Unity 端做 resize / 质量压缩，直接转发。
            return jpegBytes;
        }

        public byte[] GetLatestJpegCopy()
        {
            lock (_jpegLock)
            {
                return _latestJpeg != null ? (byte[])_latestJpeg.Clone() : null;
            }
        }

        public bool TryGetLatestResult(out LightGlueResult result)
        {
            if (_resultReceiver != null)
            {
                return _resultReceiver.TryDequeueResult(out result);
            }
            result = default(LightGlueResult);
            return false;
        }

        public void Dispose()
        {
            Stop();
        }
    }
}