poe2-bot/src/Poe2Trade.Screen/PythonDetectBridge.cs

namespace Poe2Trade.Screen;

using System.Diagnostics;
using System.Text.Json;
using System.Text.Json.Serialization;
using OpenCvSharp;
using Serilog;

/// <summary>
/// Manages a persistent Python subprocess for YOLO object detection.
/// Lazy-starts on first request; reuses the process for subsequent calls.
/// Same stdin/stdout JSON-per-line protocol as PythonOcrBridge.
/// </summary>
class PythonDetectBridge : IDisposable
{
    private static readonly JsonSerializerOptions JsonOptions = new()
    {
        PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
        DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
    };

    private Process? _proc;
    private readonly string _daemonScript;
    private readonly string _pythonExe;
    private readonly object _lock = new();

    public PythonDetectBridge()
    {
        _daemonScript = Path.GetFullPath(Path.Combine("tools", "python-detect", "daemon.py"));

        var venvPython = Path.GetFullPath(Path.Combine("tools", "python-detect", ".venv", "Scripts", "python.exe"));
        _pythonExe = File.Exists(venvPython) ? venvPython : "python";
    }

    /// <summary>
    /// Run YOLO detection on a BGR Mat. Returns parsed detection results.
    /// </summary>
    public DetectResponse Detect(Mat bgrMat, float conf = 0.3f, float iou = 0.45f, int imgsz = 640, string? model = null)
    {
        EnsureRunning();

        var imageBytes = bgrMat.ToBytes(".jpg", [(int)ImwriteFlags.JpegQuality, 95]);
        var imageBase64 = Convert.ToBase64String(imageBytes);

        var req = new Dictionary<string, object?>
        {
            ["cmd"] = "detect",
            ["imageBase64"] = imageBase64,
            ["conf"] = conf,
            ["iou"] = iou,
            ["imgsz"] = imgsz,
            ["model"] = model,
        };

        return SendRequest(req);
    }

    private DetectResponse SendRequest(object req)
    {
        var json = JsonSerializer.Serialize(req, JsonOptions);

        string responseLine;
        lock (_lock)
        {
            _proc!.StandardInput.WriteLine(json);
            _proc.StandardInput.Flush();
            responseLine = _proc.StandardOutput.ReadLine()
                ?? throw new Exception("Python detect daemon returned null");
        }

        var resp = JsonSerializer.Deserialize<PythonDetectResponse>(responseLine, JsonOptions);
        if (resp == null)
            throw new Exception("Failed to parse Python detect response");
        if (!resp.Ok)
            throw new Exception(resp.Error ?? "Python detect failed");

        return new DetectResponse
        {
            Count = resp.Count,
            InferenceMs = resp.InferenceMs,
            Detections = resp.Detections ?? [],
        };
    }

    private void EnsureRunning()
    {
        if (_proc != null && !_proc.HasExited)
            return;

        _proc?.Dispose();
        _proc = null;

        if (!File.Exists(_daemonScript))
            throw new Exception($"Python detect daemon not found at {_daemonScript}");

        Log.Information("Spawning Python detect daemon: {Python} {Script}", _pythonExe, _daemonScript);

        var proc = new Process
        {
            StartInfo = new ProcessStartInfo
            {
                FileName = _pythonExe,
                Arguments = $"\"{_daemonScript}\"",
                UseShellExecute = false,
                RedirectStandardInput = true,
                RedirectStandardOutput = true,
                RedirectStandardError = true,
                CreateNoWindow = true,
            }
        };

        proc.ErrorDataReceived += (_, e) =>
        {
            if (!string.IsNullOrEmpty(e.Data))
                Log.Debug("[python-detect] {Line}", e.Data);
        };

        try
        {
            proc.Start();
            proc.BeginErrorReadLine();

            // Wait for ready signal (up to 60s for CUDA warmup)
            var readyTask = Task.Run(() => proc.StandardOutput.ReadLine());
            if (!readyTask.Wait(TimeSpan.FromSeconds(60)))
                throw new Exception("Python detect daemon timed out waiting for ready signal");

            var readyLine = readyTask.Result;
            if (readyLine == null)
                throw new Exception("Python detect daemon exited before ready signal");

            var ready = JsonSerializer.Deserialize<PythonDetectResponse>(readyLine, JsonOptions);
            if (ready?.Ready != true)
                throw new Exception($"Python detect daemon did not send ready signal: {readyLine}");
        }
        catch
        {
            try { if (!proc.HasExited) proc.Kill(); } catch { /* best effort */ }
            proc.Dispose();
            throw;
        }

        _proc = proc;
        Log.Information("Python detect daemon ready");
    }

    public void Dispose()
    {
        if (_proc != null && !_proc.HasExited)
        {
            try
            {
                _proc.StandardInput.Close();
                _proc.WaitForExit(3000);
                if (!_proc.HasExited) _proc.Kill();
            }
            catch { /* ignore */ }
        }
        _proc?.Dispose();
        _proc = null;
    }

    // -- Response types --

    public class DetectResponse
    {
        public int Count { get; set; }
        public float InferenceMs { get; set; }
        public List<Detection> Detections { get; set; } = [];
    }

    public class Detection
    {
        [JsonPropertyName("class")]
        public string ClassName { get; set; } = "";

        public int ClassId { get; set; }
        public float Confidence { get; set; }
        public int X { get; set; }
        public int Y { get; set; }
        public int Width { get; set; }
        public int Height { get; set; }
        public int Cx { get; set; }
        public int Cy { get; set; }
    }

    private class PythonDetectResponse
    {
        public bool Ok { get; set; }
        public bool? Ready { get; set; }
        public string? Error { get; set; }
        public int Count { get; set; }
        public float InferenceMs { get; set; }
        public List<Detection>? Detections { get; set; }
    }
}