Robovoice/Robovoice.App/MainForm.cs

using NAudio.Wave;
using Robovoice.App;
using Robovoice.Core;
using Robovoice.Core.Voices;
using Robovoice.Stt.Tcp;
using Robovoice.Tts.LibPiper;
using System.Diagnostics;

namespace Robovoice.App;

internal sealed partial class MainForm : Form
{
    private readonly string _voicesDir = AppConfig.VoicesDir;
    private readonly string _espeakDataPath;
    private readonly AppConfig _config;

    private LibPiperTtsEngine? _tts;
    private AudioOutput? _audioOutput;
    private TcpSttSource? _sttSource;
    private VoicePipeline? _pipeline;
    private PttHotkey? _pttHotkey;
    private NotifyIcon? _trayIcon;
    private bool _trayInit;
    private bool _locked;

    public MainForm()
    {
        InitializeComponent();
        _espeakDataPath = Path.Combine(AppContext.BaseDirectory, "espeak-ng-data");
        _config = AppConfig.Load();
        Directory.CreateDirectory(AppConfig.AppDataDir);
        Directory.CreateDirectory(_voicesDir);
        Text = $"Robovoice {AppVersion}";
        Load += OnLoad;
        FormClosing += OnFormClosing;
    }

    private async void OnLoad(object? sender, EventArgs e)
    {
        try
        {
            ActiveControl = txtLog;
            PopulateVoices();
            PopulateOutputDevices();

            _pttKey = (Keys)_config.PttKey;
            if (_pttKey == Keys.None)
                _pttKey = Keys.F8;
            txtPttKey.Text = KeyToDisplayString(_pttKey);

            if (!string.IsNullOrEmpty(_config.Voice) && cmbVoice.Items.Contains(_config.Voice))
                cmbVoice.SelectedItem = _config.Voice;
            else if (cmbVoice.Items.Count > 0)
                cmbVoice.SelectedIndex = 0;

            if (!string.IsNullOrEmpty(_config.OutputDevice) && cmbOutput.Items.Contains(_config.OutputDevice))
                cmbOutput.SelectedItem = _config.OutputDevice;
            else
                AutoSelectCableOutput();

            trkNoise.Value = _config.NoiseScale;
            trkSpeed.Value = _config.LengthScale;
            trkNoiseW.Value = _config.NoiseWScale;
            OnSliderScroll(null, EventArgs.Empty);

            chkMinimizeToTray.Checked = _config.MinimizeToTray;
            txtSttEndpoint.Text = _config.SttEndpoint;

            btnManageVoices.Click += OnManageVoices;
            btnClearLog.Click += (_, _) => txtLog.Clear();
            btnLock.Click += OnLockToggle;
            txtPttKey.Enter += OnPttKeyFocus;
            txtPttKey.KeyDown += OnPttKeyDown;
            txtSttEndpoint.Leave += OnSttEndpointChanged;

            trkNoise.Scroll += OnSliderScroll;
            trkSpeed.Scroll += OnSliderScroll;
            trkNoiseW.Scroll += OnSliderScroll;
            trkNoise.MouseUp += OnSliderReleased;
            trkSpeed.MouseUp += OnSliderReleased;
            trkNoiseW.MouseUp += OnSliderReleased;

            chkMinimizeToTray.CheckedChanged += (_, _) => SaveConfig();

            Resize += OnResize;

            SetupTray();
        }
        catch (Exception ex)
        {
            Debug.WriteLine($"OnLoad failed: {ex}");
        }
    }

    private Keys _pttKey = Keys.F8;
    private bool _capturingPttKey;
    private PttHotkey? _captureHook;

    private void OnPttKeyFocus(object? sender, EventArgs e)
    {
        _capturingPttKey = true;
        txtPttKey.Text = "Press a key...";
        txtPttKey.BackColor = Color.LightYellow;

        _captureHook?.Dispose();
        _captureHook = new PttHotkey();
        _captureHook.CaptureKeyPressed += OnCaptureKey;
        _captureHook.InstallCaptureHook();
    }

    private void OnCaptureKey(Keys key)
    {
        if (!_capturingPttKey) return;

        if (key == Keys.Escape)
        {
            CancelCapture();
            return;
        }

        if (key is Keys.ShiftKey or Keys.Menu or Keys.LWin or Keys.RWin)
            return;

        _pttKey = key;
        _capturingPttKey = false;
        txtPttKey.Text = KeyToDisplayString(_pttKey);
        txtPttKey.BackColor = SystemColors.Window;
        _captureHook?.Dispose();
        _captureHook = null;
        SetupHotkey();
        SaveConfig();
    }

    private void CancelCapture()
    {
        _capturingPttKey = false;
        txtPttKey.Text = KeyToDisplayString(_pttKey);
        txtPttKey.BackColor = SystemColors.Window;
        _captureHook?.Dispose();
        _captureHook = null;
    }

    private void OnPttKeyDown(object? sender, KeyEventArgs e)
    {
        if (_capturingPttKey && e.KeyCode == Keys.Escape)
        {
            CancelCapture();
            e.SuppressKeyPress = true;
        }
    }

    private static string KeyToDisplayString(Keys key)
    {
        return key switch
        {
            >= Keys.F1 and <= Keys.F12 => key.ToString(),
            Keys.RControlKey => "Right Ctrl",
            Keys.LControlKey => "Left Ctrl",
            Keys.Space => "Space",
            Keys.LButton => "Mouse Left",
            Keys.RButton => "Mouse Right",
            Keys.MButton => "Mouse Middle",
            Keys.XButton1 => "Mouse X1",
            Keys.XButton2 => "Mouse X2",
            Keys.Oemtilde => "`",
            Keys.CapsLock => "CapsLock",
            Keys.NumLock => "NumLock",
            Keys.Scroll => "ScrollLock",
            Keys.Pause => "Pause",
            Keys.Insert => "Insert",
            Keys.Delete => "Delete",
            Keys.Home => "Home",
            Keys.End => "End",
            Keys.PageUp => "PageUp",
            Keys.PageDown => "PageDown",
            _ => key.ToString(),
        };
    }

    private void PopulateVoices()
    {
        cmbVoice.Items.Clear();
        if (!Directory.Exists(_voicesDir)) return;

        foreach (var onnx in Directory.GetFiles(_voicesDir, "*.onnx"))
        {
            string name = Path.GetFileNameWithoutExtension(onnx);
            cmbVoice.Items.Add(name);
        }
    }

    private void PopulateOutputDevices()
    {
        cmbOutput.Items.Clear();
        foreach (var (index, name) in AudioOutput.GetDevices())
        {
            cmbOutput.Items.Add(name);
        }
    }

    private void AutoSelectCableOutput()
    {
        for (int i = 0; i < cmbOutput.Items.Count; i++)
        {
            if (cmbOutput.Items[i] is string s && s.Contains("CABLE", StringComparison.OrdinalIgnoreCase))
            {
                cmbOutput.SelectedIndex = i;
                return;
            }
        }
        if (cmbOutput.Items.Count > 0)
            cmbOutput.SelectedIndex = 0;
    }

    // ─── Lock / Release ────────────────────────────────────────────────────

    private async void OnLockToggle(object? sender, EventArgs e)
    {
        if (_locked)
        {
            ReleaseModel();
        }
        else
        {
            await LockModelAsync();
        }
    }

    private async Task LockModelAsync()
    {
        if (cmbVoice.SelectedItem is not string voiceName)
        {
            Log("No voice selected.");
            return;
        }

        string modelPath = Path.Combine(_voicesDir, voiceName + ".onnx");
        if (!File.Exists(modelPath))
        {
            Log($"Model not found: {modelPath}");
            return;
        }

        if (!Directory.Exists(_espeakDataPath))
        {
            Log($"espeak-ng-data not found at: {_espeakDataPath}");
            return;
        }

        btnLock.Enabled = false;
        Log($"Loading voice: {voiceName}...");

        try
        {
            _tts = new LibPiperTtsEngine(
                modelPath,
                _espeakDataPath,
                noiseScale: trkNoise.Value / 1000.0f,
                lengthScale: trkSpeed.Value / 100.0f,
                noiseWScale: trkNoiseW.Value / 1000.0f);

            await Task.Run(() => _tts.InitializeAsync());
            Log($"TTS ready (sample rate: {_tts.SampleRate} Hz)");

            _audioOutput = new AudioOutput();

            _sttSource = new TcpSttSource { Endpoint = txtSttEndpoint.Text, Log = Log };
            _sttSource.TranscriptReceived += OnTranscript;
            await _sttSource.StartAsync();

            _pipeline = new VoicePipeline(_tts, _audioOutput, Log)
            {
                OutputDeviceName = cmbOutput.SelectedItem as string ?? string.Empty,
            };
            _pipeline.Start();

            _locked = true;
            SetControlsLocked(true);
            btnLock.Text = "Release Model";
            btnLock.Enabled = true;

            SetupHotkey();
            Log("Model locked. Press PTT to talk.");
        }
        catch (Exception ex)
        {
            Log($"Lock failed: {ex.Message}");
            btnLock.Enabled = true;
            _tts?.DisposeAsync().AsTask().Wait();
            _tts = null;
            _audioOutput?.Dispose();
            _audioOutput = null;
            if (_sttSource is not null)
            {
                _sttSource.TranscriptReceived -= OnTranscript;
                _sttSource.DisposeAsync().AsTask().Wait();
                _sttSource = null;
            }
        }
    }

    private void ReleaseModel()
    {
        if (!_locked) return;

        _pttHotkey?.Dispose();
        _pttHotkey = null;

        if (_sttSource is not null)
        {
            _sttSource.TranscriptReceived -= OnTranscript;
            _sttSource.DisposeAsync().AsTask().Wait();
            _sttSource = null;
        }

        _pipeline?.Dispose();
        _pipeline = null;

        // VoicePipeline.Dispose calls _tts.DisposeAsync and _audioOutput.Dispose
        _tts = null;
        _audioOutput = null;

        _locked = false;
        SetControlsLocked(false);
        btnLock.Text = "Lock Model";
        Log("Model released.");
    }

    private void SetControlsLocked(bool locked)
    {
        cmbVoice.Enabled = !locked;
        cmbOutput.Enabled = !locked;
        trkNoise.Enabled = !locked;
        trkSpeed.Enabled = !locked;
        trkNoiseW.Enabled = !locked;
        txtSttEndpoint.Enabled = !locked;
        btnManageVoices.Enabled = !locked;
    }

    // ─── PTT ────────────────────────────────────────────────────────────────

    private void SetupHotkey()
    {
        _pttHotkey?.Dispose();
        _pttHotkey = new PttHotkey { Key = _pttKey };
        _pttHotkey.Pressed += OnPttPressed;
        _pttHotkey.Released += OnPttReleased;
        _pttHotkey.Install();
        Log($"PTT hotkey installed: {_pttHotkey.Key}");
    }

    private void OnPttPressed(object? sender, EventArgs e)
    {
        _pipeline?.OnPttPressed();
        _sttSource?.SendOn();
    }

    private void OnPttReleased(object? sender, EventArgs e)
    {
        _sttSource?.SendOff();
        _pipeline?.OnPttReleased();
    }

    // ─── STT transcript → pipeline ──────────────────────────────────────────

    private void OnTranscript(object? sender, TranscriptEventArgs e)
    {
        if (_pipeline is null || !_locked) return;

        var msg = e.Message;
        if (msg.Type == TranscriptType.Partial)
        {
            if (!string.IsNullOrWhiteSpace(msg.Text))
            {
                Log($"SEGMENT: \"{msg.Text}\"");
                _pipeline.EnqueueSegment(msg.Text);
            }
        }
        else
        {
            if (!string.IsNullOrWhiteSpace(msg.Text))
            {
                Log($"FINAL: \"{msg.Text}\"");
                _pipeline.EnqueueFinal(msg.Text);
            }
            else
            {
                Log("FINAL: (empty)");
                _pipeline.EnqueueFinal("");
            }
        }
    }

    // ─── Voice manager ─────────────────────────────────────────────────────

    private void OnManageVoices(object? sender, EventArgs e)
    {
        if (_locked)
        {
            Log("Release the model before managing voices.");
            return;
        }

        string currentVoice = cmbVoice.SelectedItem as string ?? string.Empty;
        using var dlg = new VoiceManagerForm(_voicesDir, currentVoice);
        dlg.ShowDialog(this);

        PopulateVoices();

        if (!string.IsNullOrEmpty(dlg.SelectedVoiceKey) &&
            cmbVoice.Items.Contains(dlg.SelectedVoiceKey))
        {
            cmbVoice.SelectedItem = dlg.SelectedVoiceKey;
        }
        else if (cmbVoice.Items.Count > 0)
        {
            cmbVoice.SelectedIndex = 0;
        }

        SaveConfig();
    }

    // ─── Slider / endpoint ─────────────────────────────────────────────────

    private void OnSliderScroll(object? sender, EventArgs e)
    {
        lblNoiseVal.Text = $"{trkNoise.Value / 1000.0:F3}";
        lblSpeedVal.Text = $"{trkSpeed.Value / 100.0:F2}";
        lblNoiseWVal.Text = $"{trkNoiseW.Value / 1000.0:F3}";
    }

    private void OnSliderReleased(object? sender, MouseEventArgs e)
    {
        SaveConfig();
    }

    private void OnSttEndpointChanged(object? sender, EventArgs e)
    {
        _config.SttEndpoint = txtSttEndpoint.Text;
        SaveConfig();
    }

    // ─── Tray ──────────────────────────────────────────────────────────────

    private void SetupTray()
    {
        if (_trayInit) return;
        _trayInit = true;

        var micIcon = IconExtractor.TryExtractMicrophone();
        if (micIcon is not null)
        {
            _trayIcon = new NotifyIcon
            {
                Icon = micIcon,
                Text = $"Robovoice {AppVersion}",
                Visible = true,
            };
            Icon = micIcon;
        }
        else
        {
            _trayIcon = new NotifyIcon
            {
                Icon = SystemIcons.Application,
                Text = $"Robovoice {AppVersion}",
                Visible = true,
            };
        }

        var menu = new ContextMenuStrip();
        menu.Items.Add("Show", null, (_, _) => ShowWindow());
        menu.Items.Add("Exit", null, (_, _) =>
        {
            _trayIcon.Visible = false;
            Application.Exit();
        });
        _trayIcon.ContextMenuStrip = menu;
        _trayIcon.DoubleClick += (_, _) => ShowWindow();
    }

    private static string AppVersion =>
        typeof(MainForm).Assembly.GetName().Version?.ToString() ?? "0.0";

    private void OnResize(object? sender, EventArgs e)
    {
        if (WindowState == FormWindowState.Minimized && chkMinimizeToTray.Checked)
        {
            Hide();
            WindowState = FormWindowState.Normal;
        }
    }

    private void ShowWindow()
    {
        Show();
        WindowState = FormWindowState.Normal;
        Activate();
    }

    // ─── Config ────────────────────────────────────────────────────────────

    private void SaveConfig()
    {
        _config.Voice = cmbVoice.SelectedItem as string ?? string.Empty;
        _config.PttKey = (int)_pttKey;
        _config.OutputDevice = cmbOutput.SelectedItem as string ?? string.Empty;
        _config.NoiseScale = trkNoise.Value;
        _config.LengthScale = trkSpeed.Value;
        _config.NoiseWScale = trkNoiseW.Value;
        _config.MinimizeToTray = chkMinimizeToTray.Checked;
        _config.SttEndpoint = txtSttEndpoint.Text;
        _config.Save();
    }

    // ─── Logging ───────────────────────────────────────────────────────────

    private void Log(string message)
    {
        if (IsDisposed) return;
        if (InvokeRequired)
        {
            BeginInvoke(() => Log(message));
            return;
        }

        string timestamp = DateTime.Now.ToString("HH:mm:ss");
        txtLog.AppendText($"[{timestamp}] {message}\n");
        txtLog.ScrollToCaret();
    }

    // ─── Shutdown ──────────────────────────────────────────────────────────

    private void OnFormClosing(object? sender, FormClosingEventArgs e)
    {
        _pttHotkey?.Dispose();
        _trayIcon!.Visible = false;
        ReleaseModel();
        SaveConfig();
    }
}