HawkeyeVision/VisionBuilder.UI.RaspberryAcquisition/ViewModels/MainWindowViewModel.cs

using System.Globalization;
using Avalonia.Threading;
using CommunityToolkit.Mvvm.ComponentModel;
using CommunityToolkit.Mvvm.Input;
using OpenCvSharp;
using VisionBuilder.UI.RaspberryAcquisition.Models;
using VisionBuilder.UI.RaspberryAcquisition.Services;

namespace VisionBuilder.UI.RaspberryAcquisition.ViewModels;

public partial class MainWindowViewModel : ObservableObject, IDisposable
{
    private const string CalibrationFile = "calibration.json";

    private ConfigurableLibcameraSource? _source;
    private CancellationTokenSource? _previewCts;
    private Mat? _lastFrame;
    private Mat? _cameraMatrix;
    private Mat? _distCoeffs;

    [ObservableProperty] private Avalonia.Media.Imaging.Bitmap? _previewImage;
    [ObservableProperty] private string _statusText = "Stopped";
    [ObservableProperty] private int _frameCount;
    [ObservableProperty] private bool _isRunning;
    [ObservableProperty] private bool _isCalibrated;

    [ObservableProperty] private string _width = "2028";
    [ObservableProperty] private string _height = "1520";
    [ObservableProperty] private string _shutterSpeed = "10000";
    [ObservableProperty] private string _framerate = "2";
    [ObservableProperty] private string _awbGainRed = "3.86";
    [ObservableProperty] private string _awbGainBlue = "1.46";
    [ObservableProperty] private string _saveDirectory = "/home/oem/captures";
    [ObservableProperty] private string _calibrationDirectory = "/home/oem/captures";

    public MainWindowViewModel()
    {
        TryLoadCalibration();
    }

    private void TryLoadCalibration()
    {
        if (!File.Exists(CalibrationFile)) return;

        try
        {
            (_cameraMatrix, _distCoeffs) = CameraCalibrationService.LoadCalibration(CalibrationFile);
            IsCalibrated = true;
            StatusText = "Calibration loaded";
        }
        catch (Exception ex)
        {
            StatusText = $"Calibration load error: {ex.Message}";
        }
    }

    [RelayCommand]
    private void StartCapture()
    {
        if (IsRunning) return;

        var settings = BuildSettings();
        if (settings == null) return;

        _source = new ConfigurableLibcameraSource(settings);
        _source.StatusChanged += status =>
            Dispatcher.UIThread.Post(() => StatusText = status);

        _source.Start();
        IsRunning = true;

        _previewCts = new CancellationTokenSource();
        _ = PreviewLoopAsync(_previewCts.Token);
    }

    [RelayCommand]
    private void StopCapture()
    {
        if (!IsRunning) return;

        _previewCts?.Cancel();
        _source?.Stop();
        _source?.Dispose();
        _source = null;
        IsRunning = false;
        StatusText = "Stopped";
    }

    [RelayCommand]
    private void ApplySettings()
    {
        if (!IsRunning || _source == null) return;

        var settings = BuildSettings();
        if (settings == null) return;

        _source.UpdateSettings(settings);
    }

    [RelayCommand]
    private void SaveImage()
    {
        if (_lastFrame == null || _lastFrame.Empty()) return;

        Directory.CreateDirectory(SaveDirectory);
        var filename = $"capture_{DateTime.Now:yyyyMMdd_HHmmss}.png";
        var path = Path.Combine(SaveDirectory, filename);
        Cv2.ImWrite(path, _lastFrame);
        StatusText = $"Saved: {filename}";
    }

    [RelayCommand]
    private void CalibrateCamera()
    {
        try
        {
            StatusText = "Calibrating...";
            var data = CameraCalibrationService.Calibrate(CalibrationDirectory, new Size(9, 6));
            CameraCalibrationService.SaveCalibration(data, CalibrationFile);

            _cameraMatrix?.Dispose();
            _distCoeffs?.Dispose();
            (_cameraMatrix, _distCoeffs) = CameraCalibrationService.LoadCalibration(CalibrationFile);

            IsCalibrated = true;
            StatusText = $"Calibrated (RMS: {data.RmsError:F4})";
        }
        catch (Exception ex)
        {
            StatusText = $"Calibration error: {ex.Message}";
        }
    }

    private async Task PreviewLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested && _source != null)
        {
            try
            {
                var frame = await _source.GetImage(ct);

                if (_cameraMatrix != null && _distCoeffs != null)
                {
                    var corrected = new Mat();
                    Cv2.Undistort(frame, corrected, _cameraMatrix, _distCoeffs);
                    frame.Dispose();
                    frame = corrected;
                }

                _lastFrame?.Dispose();
                _lastFrame = frame;

                var bitmap = ImageConverter.MatToAvaloniaBitmap(frame);
                await Dispatcher.UIThread.InvokeAsync(() =>
                {
                    var old = PreviewImage;
                    PreviewImage = bitmap;
                    FrameCount = _source?.FrameCount ?? 0;
                    old?.Dispose();
                });
            }
            catch (OperationCanceledException)
            {
                break;
            }
            catch (Exception ex)
            {
                await Dispatcher.UIThread.InvokeAsync(() =>
                    StatusText = $"Preview error: {ex.Message}");
                break;
            }
        }
    }

    private CaptureSettings? BuildSettings()
    {
        if (!int.TryParse(Width, out var w) ||
            !int.TryParse(Height, out var h) ||
            !int.TryParse(ShutterSpeed, out var s) ||
            !int.TryParse(Framerate, out var f) ||
            !double.TryParse(AwbGainRed, CultureInfo.InvariantCulture, out var r) ||
            !double.TryParse(AwbGainBlue, CultureInfo.InvariantCulture, out var b))
        {
            StatusText = "Error: invalid settings values";
            return null;
        }

        return new CaptureSettings
        {
            Width = w,
            Height = h,
            ShutterSpeed = s,
            Framerate = f,
            AwbGainRed = r,
            AwbGainBlue = b
        };
    }

    public void Dispose()
    {
        StopCapture();
        _lastFrame?.Dispose();
        _cameraMatrix?.Dispose();
        _distCoeffs?.Dispose();
    }
}