smart-bow-infrared/Assets/LightGlue/Scripts/Roma/Bridge/RomaHardwareToPythonBridge.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.Roma.Bridge
{
    /// <summary>
    /// Roma 版本的桥接层：目标是最大化复用旧 HardwareToPythonUdpBridge 的“能力形态”，但核心差异是：
    /// - 硬件图像不再先到 Unity；Python 直接监听硬件 FrameHeader 图传端口（例如 9000）
    ///
    /// 保留/提供的能力：
    /// - Unity -> Hardware：通过 RomaWifiCameraControlClient 发送 config/stream_on/stream_off
    /// - Unity -> Python：控制指令 n/s/r（pythonControlPort=12349）
    /// - Unity -> Python：可选发送 Unity GameView JPEG（单包 UDP，unityRefPort=12347），用于 r 参考图更新
    /// - Python -> Unity：结果接收 UDPResultReceiver（12348）
    /// - Python 进程：PythonProcessController 启动/停止
    /// - Unity 显示：可选启动 RomaForwardedImageViewer（12366）
    /// </summary>
    public sealed class RomaHardwareToPythonBridge : MonoBehaviour
    {
        [Header("配置管理")]
        public bool autoLoadNetworkConfig = true;

        [Header("Discovery + WiFi Control (Unity -> Hardware)")]
        public RomaDeviceDiscoveryListener discovery;
        public RomaWifiCameraControlClient wifiControl;

        [Header("ESP32 Hardware Control (Unity -> ESP32, 0x40)")]
        [Tooltip("是否启用 ESP32 硬件 0x40 参数下发（供 ImageTransmissionUIController 调用）。")]
        public bool enableEsp32HardwareControl = true;

        [Tooltip("ESP32 设备IP（用于发送 0x40 参数帧）。默认从 NetworkConfig.romaEsp32DeviceIp 加载。")]
        public string esp32DeviceIp = "";

        [Tooltip("ESP32 设备端口（用于发送 0x40 参数帧），默认 8888。")]
        public int esp32DevicePort = 8888;

        [Header("Python Process")]
        public PythonProcessController pythonProcessController;

        [Header("Python Control (Unity -> Python)")]
        public bool enablePythonControl = true;
        public string pythonIp = "127.0.0.1";
        public int pythonControlPort = 12349;

        [Header("Unity GameView Reference (optional)")]
        public bool enableUnityGameViewReference = true;
        public int unityRefPort = 12347;
        public Camera gameViewCamera;
        public int referenceCaptureWidth = 320;
        public int referenceCaptureHeight = 240;
        [Range(1, 100)] public int referenceJpegQuality = 80;

        [Header("Reference Source")]
        public ReferenceSource referenceSource = ReferenceSource.CameraOnly;
        [Range(0.01f, 60f)] public float referenceRefreshIntervalSeconds = 5f;

        [Header("Python Result Receiver (Python -> Unity)")]
        public bool enableResultReceiver = true;
        [Tooltip("是否在 OnEnable 自动启动结果接收（手动流程建议关闭，避免端口冲突）。")]
        public bool autoStartResultReceiverOnEnable = false;
        public string pythonResultBindIp = "127.0.0.1";
        public int pythonResultPort = 12348;
        public int maxResultQueueSize = 10;

        [Header("Debug")]
        [Tooltip("收到结果时是否打印日志（限流）。")]
        public bool logResults = true;
        [Min(0.1f)]
        public float logResultIntervalSeconds = 1.0f;

        [Header("Forwarded Image Viewer (Python -> Unity image)")]
        public RomaForwardedImageViewer forwardedViewer;

        public enum ReferenceSource
        {
            CameraOnly,     // 仅手动触发 n
            ScreenCapture,  // 周期性触发 s（Python 截屏）
            GameView        // 周期性发送 GameView JPEG + 触发 r
        }

        private UDPResultReceiver _resultReceiver;
        public event Action<LightGlueResult> OnResultUpdated;
        private LightGlueResult _latest;
        private bool _hasLatest;
        public bool HasLatestResult => _hasLatest;
        public LightGlueResult LatestResult => _latest;

        private UdpClient _pyControlClient;
        private IPEndPoint _pyControlEndpoint;

        private UdpClient _unityRefSender;
        private IPEndPoint _unityRefEndpoint;

        private UdpClient _esp32ControlClient;
        private IPEndPoint _esp32ControlEndpoint;

        private float _lastReferenceRefreshTime;
        private float _lastResultLogTime;

        private void OnEnable()
        {
            Application.runInBackground = true;
            _lastReferenceRefreshTime = Time.unscaledTime;

            ReferenceImageResizeService.EnsureInitialized();
            referenceCaptureWidth = ReferenceImageResizeService.Width;
            referenceCaptureHeight = ReferenceImageResizeService.Height;
            ReferenceImageResizeService.OnResizeChanged += OnResizeChanged;

            if (autoLoadNetworkConfig)
                LoadNetworkConfig();

            if (enablePythonControl)
                InitPythonControl();

            if (enableUnityGameViewReference)
                InitUnityRefSender();

            if (enableResultReceiver && autoStartResultReceiverOnEnable)
                StartResultReceiver();
        }

        private void OnDisable()
        {
            StopResultReceiver();
            try { _pyControlClient?.Close(); } catch { /* ignore */ }
            _pyControlClient = null;
            try { _unityRefSender?.Close(); } catch { /* ignore */ }
            _unityRefSender = null;
            try { _esp32ControlClient?.Close(); } catch { /* ignore */ }
            _esp32ControlClient = null;

            ReferenceImageResizeService.OnResizeChanged -= OnResizeChanged;
        }

        private void OnResizeChanged(int w, int h)
        {
            referenceCaptureWidth = w;
            referenceCaptureHeight = h;
        }

        private void LoadNetworkConfig()
        {
            try
            {
                NetworkConfig cfg = NetworkConfigManager.LoadConfig();
                if (cfg == null || !cfg.Validate()) return;

                pythonIp = cfg.pythonIp;
                pythonResultBindIp = cfg.pythonResultBindIp;
                pythonResultPort = cfg.pythonResultPort;

                // Roma WiFi 端口：用于下发给硬件的 target_port（Python监听硬件图像）
                if (wifiControl != null)
                {
                    wifiControl.targetPort = cfg.romaWifiListenPort;
                }

                // Roma forward viewer 端口
                if (forwardedViewer != null)
                    forwardedViewer.forwardPort = cfg.romaForwardPort;

                // Roma 控制端口
                pythonControlPort = cfg.romaControlPort;

                // UnityRefPort（默认复用 pythonPort=12347）
                unityRefPort = cfg.pythonPort;

                // ESP32 控制端点（仅用于 0x40 下发）
                esp32DeviceIp = cfg.romaEsp32DeviceIp;
                esp32DevicePort = cfg.romaEsp32DevicePort;
                InitEsp32Control();
            }
            catch (Exception ex)
            {
                Debug.LogWarning($"[RomaBridge] Load config failed: {ex.Message}");
            }
        }

        private void InitEsp32Control()
        {
            if (!enableEsp32HardwareControl) return;
            if (string.IsNullOrWhiteSpace(esp32DeviceIp)) return;
            if (esp32DevicePort <= 0 || esp32DevicePort > 65535) return;

            try
            {
                _esp32ControlEndpoint = new IPEndPoint(IPAddress.Parse(esp32DeviceIp.Trim()), esp32DevicePort);
                _esp32ControlClient ??= new UdpClient();
            }
            catch (Exception ex)
            {
                Debug.LogWarning($"[RomaBridge] Init ESP32 control failed: {ex.Message}");
                _esp32ControlEndpoint = null;
                try { _esp32ControlClient?.Close(); } catch { /* ignore */ }
                _esp32ControlClient = null;
            }
        }

        /// <summary>
        /// 供 UI 调用：将 ImageTransmissionConfig 以 0x40 参数帧下发给 ESP32。
        /// </summary>
        public void ApplyEsp32HardwareConfig(ImageTransmissionConfig config)
        {
            if (!enableEsp32HardwareControl)
            {
                Debug.LogWarning("[RomaBridge] ESP32 hardware control disabled, skip 0x40 apply.");
                return;
            }

            if (_esp32ControlClient == null || _esp32ControlEndpoint == null)
                InitEsp32Control();
            if (_esp32ControlClient == null || _esp32ControlEndpoint == null)
            {
                Debug.LogWarning("[RomaBridge] ESP32 control endpoint not ready (device ip/port empty?), cannot send 0x40.");
                return;
            }

            try
            {
                byte[] frame = CoreHardwareControlProtocol.BuildParameterSettingFrame(config);
                if (frame == null || frame.Length == 0) return;
                _esp32ControlClient.Send(frame, frame.Length, _esp32ControlEndpoint);
                Debug.Log($"[RomaBridge] Sent 0x40 frame -> {_esp32ControlEndpoint.Address}:{_esp32ControlEndpoint.Port}");
            }
            catch (Exception ex)
            {
                Debug.LogWarning($"[RomaBridge] Send 0x40 to ESP32 failed: {ex.Message}");
            }
        }

        private void InitPythonControl()
        {
            try
            {
                string ipStr = string.IsNullOrWhiteSpace(pythonIp) ? "127.0.0.1" : pythonIp;
                _pyControlEndpoint = new IPEndPoint(IPAddress.Parse(ipStr), pythonControlPort);
                _pyControlClient = new UdpClient();
            }
            catch (Exception ex)
            {
                Debug.LogError($"[RomaBridge] Init python control failed: {ex.Message}");
                _pyControlClient = null;
            }
        }

        private void InitUnityRefSender()
        {
            try
            {
                string ipStr = string.IsNullOrWhiteSpace(pythonIp) ? "127.0.0.1" : pythonIp;
                _unityRefEndpoint = new IPEndPoint(IPAddress.Parse(ipStr), unityRefPort);
                _unityRefSender = new UdpClient();
            }
            catch (Exception ex)
            {
                Debug.LogError($"[RomaBridge] Init unity ref sender failed: {ex.Message}");
                _unityRefSender = null;
            }
        }

        public void StartPython()
        {
            if (pythonProcessController == null) return;
            pythonProcessController.StartPython();
        }

        public void StopPython()
        {
            if (pythonProcessController == null) return;
            pythonProcessController.StopPython();
        }

        public void StartForwardViewer()
        {
            forwardedViewer?.StartReceiver();
        }

        public void StopForwardViewer()
        {
            forwardedViewer?.StopReceiver();
        }

        public void SendControlCommand(char cmd)
        {
            if (!enablePythonControl || _pyControlClient == null || _pyControlEndpoint == null)
                return;
            try
            {
                byte[] payload = { (byte)cmd };
                _pyControlClient.Send(payload, payload.Length, _pyControlEndpoint);
            }
            catch (Exception ex)
            {
                Debug.LogWarning($"[RomaBridge] Send control '{cmd}' failed: {ex.Message}");
            }
        }

        public void SendWifiConfigToHardware(string targetIp, int targetPort)
        {
            if (wifiControl == null) return;
            wifiControl.ConfigStream(targetIp, targetPort);
        }

        public void StreamOn()
        {
            wifiControl?.StartStream();
        }

        public void StreamOff()
        {
            wifiControl?.StopStream();
        }

        private void StartResultReceiver()
        {
            if (_resultReceiver != null) return;
            try
            {
                _resultReceiver = new UDPResultReceiver(pythonResultBindIp, pythonResultPort, maxResultQueueSize);
                _resultReceiver.Start();
            }
            catch (Exception ex)
            {
                Debug.LogError($"[RomaBridge] Start result receiver failed: {ex.Message}");
                _resultReceiver = null;
            }
        }

        private void StopResultReceiver()
        {
            try { _resultReceiver?.Stop(); } catch { /* ignore */ }
            _resultReceiver = null;
            _hasLatest = false;
        }

        public void EnsureResultReceiverStarted()
        {
            if (!enableResultReceiver) return;
            if (_resultReceiver != null) return;
            StartResultReceiver();
        }

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

            bool any = false;
            while (_resultReceiver.TryDequeueResult(out var r))
            {
                result = r;
                any = true;
            }
            return any;
        }

        private void Update()
        {
            // drain results and publish
            if (_resultReceiver != null)
            {
                bool any = false;
                LightGlueResult last = default;
                while (_resultReceiver.TryDequeueResult(out var r))
                {
                    last = r;
                    any = true;
                }
                if (any)
                {
                    _latest = last;
                    _hasLatest = true;
                    OnResultUpdated?.Invoke(last);

                    if (logResults && (Time.unscaledTime - _lastResultLogTime) >= logResultIntervalSeconds)
                    {
                        _lastResultLogTime = Time.unscaledTime;
                        Debug.Log($"[RomaBridge] Result: valid={(_latest.IsValid ? 1 : 0)} matches={_latest.NumMatches} inliers={_latest.InliersRatio:F3} x={_latest.CameraPosition.x:F1} y={_latest.CameraPosition.y:F1}");
                    }
                }
            }

            if (!enablePythonControl) return;

            if (referenceSource == ReferenceSource.ScreenCapture || referenceSource == ReferenceSource.GameView)
            {
                float now = Time.unscaledTime;
                if (now - _lastReferenceRefreshTime >= referenceRefreshIntervalSeconds)
                {
                    _lastReferenceRefreshTime = now;
                    if (referenceSource == ReferenceSource.ScreenCapture)
                    {
                        SendControlCommand('s');
                    }
                    else
                    {
                        // GameView: send JPEG to python unity_ref_port, then trigger 'r'
                        if (enableUnityGameViewReference)
                        {
                            byte[] jpeg = CaptureGameViewToJpeg();
                            if (jpeg != null && jpeg.Length > 0)
                            {
                                SendUnityRefJpeg(jpeg);
                                SendControlCommand('r');
                            }
                        }
                    }
                }
            }
        }

        private void SendUnityRefJpeg(byte[] jpeg)
        {
            if (_unityRefSender == null || _unityRefEndpoint == null) return;
            try
            {
                // 单包 UDP：要求 jpeg 长度 < 65507
                _unityRefSender.Send(jpeg, jpeg.Length, _unityRefEndpoint);
            }
            catch (Exception ex)
            {
                Debug.LogWarning($"[RomaBridge] Send UnityRef JPEG failed: {ex.Message}");
            }
        }

        private byte[] CaptureGameViewToJpeg()
        {
            Camera cam = gameViewCamera != null ? gameViewCamera : Camera.main;
            if (cam == null) return null;

            int w = Mathf.Clamp(referenceCaptureWidth, 64, 4096);
            int h = Mathf.Clamp(referenceCaptureHeight, 64, 4096);

            RenderTexture rt = RenderTexture.GetTemporary(w, h, 24, RenderTextureFormat.ARGB32);
            RenderTexture prevTarget = cam.targetTexture;
            RenderTexture prevActive = RenderTexture.active;

            cam.targetTexture = rt;
            cam.Render();

            RenderTexture.active = rt;
            Texture2D tex = new Texture2D(w, h, TextureFormat.RGB24, false);
            tex.ReadPixels(new Rect(0, 0, w, h), 0, 0);
            tex.Apply();

            cam.targetTexture = prevTarget;
            RenderTexture.active = prevActive;
            RenderTexture.ReleaseTemporary(rt);

            int quality = Mathf.Clamp(referenceJpegQuality, 1, 100);
            byte[] jpeg = tex.EncodeToJPG(quality);
            Destroy(tex);
            return jpeg;
        }
    }
}