1using System;
  2using System.Diagnostics;
  3using System.IO;
  4using System.Net;
  5using System.Net.Sockets;
  6using System.Text;
  7using System.Threading;
  8
  9namespace WindowCapture.Helpers
 10{
 11    /// <summary>
 12    /// Full-text Russian correction via a WARM SAGE server (Spell/wc_spell_server.py →
 13    /// ai-forever/sage-fredt5-distilled-95m). The model is loaded ONCE in a background Python
 14    /// process; corrections then take ~0.3s instead of ~8s/call. Lazily started on first use,
 15    /// kept alive for the app's lifetime, killed on exit. Requires Python + transformers + the
 16    /// model; returns null on any failure so the caller falls back to the heuristic.
 17    /// </summary>
 18    public static class SageClient
 19    {
 20        private static readonly object gate = new object();
 21        private static Process server;
 22        private static int port;
 23        private static bool exitHookAdded;
 24
 25        private static string ScriptPath()
 26        {
 27            return Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Spell", "wc_spell_server.py");
 28        }
 29
 30        public static bool IsAvailable { get { return File.Exists(ScriptPath()); } }
 31
 32        private static int FreeTcpPort()
 33        {
 34            var l = new TcpListener(IPAddress.Loopback, 0);
 35            l.Start();
 36            int p = ((IPEndPoint)l.LocalEndpoint).Port;
 37            l.Stop();
 38            return p;
 39        }
 40
 41        private static bool Ping()
 42        {
 43            try
 44            {
 45                var r = (HttpWebRequest)WebRequest.Create("http://127.0.0.1:" + port + "/ping");
 46                r.Timeout = 1000;
 47                using (var resp = (HttpWebResponse)r.GetResponse())
 48                    return resp.StatusCode == HttpStatusCode.OK;
 49            }
 50            catch { return false; }
 51        }
 52
 53        private static bool EnsureServer()
 54        {
 55            lock (gate)
 56            {
 57                if (server != null && !server.HasExited && Ping()) return true;
 58
 59                string script = ScriptPath();
 60                if (!File.Exists(script)) return false;
 61
 62                try { if (server != null && !server.HasExited) server.Kill(); } catch { }
 63                server = null;
 64
 65                try
 66                {
 67                    port = FreeTcpPort();
 68                    var psi = new ProcessStartInfo
 69                    {
 70                        FileName = "python",
 71                        Arguments = "\"" + script + "\" " + port,
 72                        UseShellExecute = false,
 73                        CreateNoWindow = true
 74                    };
 75                    server = Process.Start(psi);
 76                    if (!exitHookAdded)
 77                    {
 78                        AppDomain.CurrentDomain.ProcessExit += delegate { Shutdown(); };
 79                        exitHookAdded = true;
 80                    }
 81                }
 82                catch { server = null; return false; }
 83
 84                // /ping only succeeds once the model has finished loading (server binds after load).
 85                for (int i = 0; i < 120; i++) // up to ~60s
 86                {
 87                    if (server == null || server.HasExited) return false;
 88                    if (Ping()) return true;
 89                    Thread.Sleep(500);
 90                }
 91                return false;
 92            }
 93        }
 94
 95        /// <summary>Ensure the shared spell server is running; return its port, or -1 on failure.
 96        /// Shared with RescoreClient (/rescore) so only ONE python process serves both SAGE
 97        /// full-text correction and the per-word context rescorer.</summary>
 98        public static int EnsurePort()
 99        {
100            return EnsureServer() ? port : -1;
101        }
102
103        /// <summary>Correct full text; null on failure (caller falls back to the heuristic).</summary>
104        public static string Correct(string text)
105        {
106            if (string.IsNullOrEmpty(text)) return text;
107            if (!File.Exists(ScriptPath())) return null;
108            if (!EnsureServer()) return null;
109            try
110            {
111                var req = (HttpWebRequest)WebRequest.Create("http://127.0.0.1:" + port + "/");
112                req.Method = "POST";
113                req.Timeout = 60000;
114                req.ContentType = "text/plain; charset=utf-8";
115                byte[] body = Encoding.UTF8.GetBytes(text);
116                req.ContentLength = body.Length;
117                using (var s = req.GetRequestStream()) s.Write(body, 0, body.Length);
118                using (var resp = (HttpWebResponse)req.GetResponse())
119                using (var rs = resp.GetResponseStream())
120                using (var sr = new StreamReader(rs, Encoding.UTF8))
121                    return sr.ReadToEnd();
122            }
123            catch { return null; }
124        }
125
126        /// <summary>Pre-load the model in the background (call when the TextAssist app starts) so the
127        /// first real correction isn't slowed by the model load.</summary>
128        public static void WarmUp()
129        {
130            if (!IsAvailable) return;
131            ThreadPool.QueueUserWorkItem(delegate { try { EnsureServer(); } catch { } });
132        }
133
134        private static void Shutdown()
135        {
136            try { lock (gate) { if (server != null && !server.HasExited) server.Kill(); } } catch { }
137        }
138    }
139}