meeting-assistant/MeetingAssistant/Screenshots/MeetingScreenshotService.cs

using System.Collections.Concurrent;
using System.Drawing;
using System.Drawing.Imaging;
using System.Runtime.InteropServices;
using MeetingAssistant.MeetingNotes;

namespace MeetingAssistant.Screenshots;

public interface IActiveWindowScreenshotCapture
{
    Task<ActiveWindowScreenshot> CaptureAsync(CancellationToken cancellationToken);
}

public sealed record ActiveWindowScreenshot(byte[] PngBytes, string? WindowTitle);

public interface IScreenshotOcrClient
{
    Task<ScreenshotOcrResult> ExtractAsync(
        string screenshotPath,
        string prompt,
        MeetingAssistantOptions options,
        CancellationToken cancellationToken);
}

public sealed record ScreenshotOcrResult(string Text, ScreenshotCropCoordinates? Crop);

public sealed record ScreenshotCropCoordinates(int X, int Y, int Width, int Height);

public interface IMeetingScreenshotService
{
    Task<MeetingScreenshotCaptureResult> CaptureAsync(
        MeetingSessionArtifacts artifacts,
        DateTimeOffset? meetingStartedAt,
        DateTimeOffset capturedAt,
        MeetingAssistantOptions options,
        CancellationToken cancellationToken);

    Task WaitForPendingOcrAsync(
        MeetingSessionArtifacts artifacts,
        TimeSpan timeout,
        CancellationToken cancellationToken);

    Task CancelPendingOcrAsync(
        MeetingSessionArtifacts artifacts,
        TimeSpan timeout,
        CancellationToken cancellationToken);
}

public sealed record MeetingScreenshotCaptureResult(
    string ScreenshotPath,
    TimeSpan MeetingTimestamp,
    bool OcrStarted);

public sealed class NoopMeetingScreenshotService : IMeetingScreenshotService
{
    public static NoopMeetingScreenshotService Instance { get; } = new();

    public Task<MeetingScreenshotCaptureResult> CaptureAsync(
        MeetingSessionArtifacts artifacts,
        DateTimeOffset? meetingStartedAt,
        DateTimeOffset capturedAt,
        MeetingAssistantOptions options,
        CancellationToken cancellationToken)
    {
        return Task.FromResult(new MeetingScreenshotCaptureResult("", TimeSpan.Zero, OcrStarted: false));
    }

    public Task WaitForPendingOcrAsync(
        MeetingSessionArtifacts artifacts,
        TimeSpan timeout,
        CancellationToken cancellationToken)
    {
        return Task.CompletedTask;
    }

    public Task CancelPendingOcrAsync(
        MeetingSessionArtifacts artifacts,
        TimeSpan timeout,
        CancellationToken cancellationToken)
    {
        return Task.CompletedTask;
    }
}

public sealed class MeetingScreenshotService : IMeetingScreenshotService
{
    private readonly IActiveWindowScreenshotCapture screenshotCapture;
    private readonly IMeetingArtifactStore artifactStore;
    private readonly IScreenshotOcrClient ocrClient;
    private readonly ILogger<MeetingScreenshotService> logger;
    private readonly ConcurrentDictionary<string, List<PendingOcrTask>> pendingOcrByContext = new(StringComparer.OrdinalIgnoreCase);
    private readonly ConcurrentDictionary<string, SemaphoreSlim> contextFileLocks = new(StringComparer.OrdinalIgnoreCase);

    public MeetingScreenshotService(
        IActiveWindowScreenshotCapture screenshotCapture,
        IMeetingArtifactStore artifactStore,
        IScreenshotOcrClient ocrClient,
        ILogger<MeetingScreenshotService> logger)
    {
        this.screenshotCapture = screenshotCapture;
        this.artifactStore = artifactStore;
        this.ocrClient = ocrClient;
        this.logger = logger;
    }

    public async Task<MeetingScreenshotCaptureResult> CaptureAsync(
        MeetingSessionArtifacts artifacts,
        DateTimeOffset? meetingStartedAt,
        DateTimeOffset capturedAt,
        MeetingAssistantOptions options,
        CancellationToken cancellationToken)
    {
        var screenshot = await screenshotCapture.CaptureAsync(cancellationToken);
        var timestamp = CalculateMeetingTimestamp(meetingStartedAt, capturedAt);
        var screenshotPath = await SaveScreenshotAsync(
            artifacts,
            screenshot.PngBytes,
            timestamp,
            capturedAt,
            options,
            cancellationToken);
        var relativePath = ToMarkdownPath(Path.GetRelativePath(
            Path.GetDirectoryName(artifacts.AssistantContextPath)!,
            screenshotPath));
        var screenshotId = Guid.NewGuid().ToString("N");
        var ocrEnabled = options.Screenshots.Ocr.Enabled;
        await artifactStore.AppendAssistantContextAsync(
            artifacts,
            CreateScreenshotMarkdown(screenshotId, timestamp, relativePath, screenshot.WindowTitle, ocrEnabled),
            cancellationToken);

        if (ocrEnabled)
        {
            var ocrCancellation = new CancellationTokenSource();
            TrackOcr(
                artifacts,
                ocrCancellation,
                RunOcrAsync(artifacts, screenshotPath, screenshotId, options, ocrCancellation.Token));
        }

        logger.LogInformation(
            "Captured meeting screenshot {ScreenshotPath} at {MeetingTimestamp}",
            screenshotPath,
            FormatTimestamp(timestamp));
        return new MeetingScreenshotCaptureResult(screenshotPath, timestamp, ocrEnabled);
    }

    public async Task WaitForPendingOcrAsync(
        MeetingSessionArtifacts artifacts,
        TimeSpan timeout,
        CancellationToken cancellationToken)
    {
        var contextKey = NormalizeContextKey(artifacts.AssistantContextPath);
        if (!pendingOcrByContext.TryGetValue(contextKey, out var tasks))
        {
            return;
        }

        PendingOcrTask[] snapshot;
        lock (tasks)
        {
            tasks.RemoveAll(task => task.Task.IsCompleted);
            snapshot = tasks.ToArray();
        }

        if (snapshot.Length == 0)
        {
            return;
        }

        try
        {
            await Task.WhenAll(snapshot.Select(task => task.Task)).WaitAsync(timeout, cancellationToken);
        }
        catch (TimeoutException)
        {
            logger.LogWarning(
                "Timed out waiting for {Count} screenshot OCR task(s) for {AssistantContextPath}",
                snapshot.Length,
                artifacts.AssistantContextPath);
        }
    }

    public async Task CancelPendingOcrAsync(
        MeetingSessionArtifacts artifacts,
        TimeSpan timeout,
        CancellationToken cancellationToken)
    {
        var contextKey = NormalizeContextKey(artifacts.AssistantContextPath);
        if (!pendingOcrByContext.TryGetValue(contextKey, out var tasks))
        {
            return;
        }

        PendingOcrTask[] snapshot;
        lock (tasks)
        {
            tasks.RemoveAll(task => task.Task.IsCompleted);
            snapshot = tasks.ToArray();
        }

        if (snapshot.Length == 0)
        {
            return;
        }

        foreach (var pending in snapshot)
        {
            await pending.Cancellation.CancelAsync();
        }

        await WaitForPendingOcrAsync(artifacts, timeout, cancellationToken);
    }

    private async Task<string> SaveScreenshotAsync(
        MeetingSessionArtifacts artifacts,
        byte[] pngBytes,
        TimeSpan timestamp,
        DateTimeOffset capturedAt,
        MeetingAssistantOptions options,
        CancellationToken cancellationToken)
    {
        var folder = ResolveAttachmentsFolder(artifacts.AssistantContextPath, options.Screenshots.AttachmentsFolder);
        Directory.CreateDirectory(folder);
        var path = Path.Combine(
            folder,
            $"{capturedAt:yyyyMMdd-HHmmss-fff}-{FormatTimestamp(timestamp).Replace(":", "", StringComparison.Ordinal)}.png");
        await File.WriteAllBytesAsync(path, pngBytes, cancellationToken);
        return path;
    }

    private async Task RunOcrAsync(
        MeetingSessionArtifacts artifacts,
        string screenshotPath,
        string screenshotId,
        MeetingAssistantOptions options,
        CancellationToken cancellationToken)
    {
        var timeout = options.Screenshots.Ocr.GetEffectiveTimeout();
        using var timeoutSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeoutSource.CancelAfter(timeout);
        try
        {
            var result = await ocrClient.ExtractAsync(
                screenshotPath,
                string.IsNullOrWhiteSpace(options.Screenshots.Ocr.Prompt)
                    ? ScreenshotOcrOptions.DefaultPrompt
                    : options.Screenshots.Ocr.Prompt,
                options,
                timeoutSource.Token);
            var cropMarkdown = await TryCreateCropMarkdownAsync(
                artifacts.AssistantContextPath,
                screenshotPath,
                result.Crop,
                timeoutSource.Token);
            await ReplaceOcrPlaceholderAsync(
                artifacts.AssistantContextPath,
                screenshotId,
                cropMarkdown +
                "### OCR" + Environment.NewLine + Environment.NewLine + result.Text.Trim(),
                CancellationToken.None);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            logger.LogInformation("Cancelled screenshot OCR for {ScreenshotPath}", screenshotPath);
        }
        catch (OperationCanceledException)
        {
            await ReplaceOcrPlaceholderAsync(
                artifacts.AssistantContextPath,
                screenshotId,
                $"_OCR timed out after {timeout:g}._",
                CancellationToken.None);
        }
        catch (Exception exception)
        {
            logger.LogWarning(exception, "Screenshot OCR failed for {ScreenshotPath}", screenshotPath);
            await ReplaceOcrPlaceholderAsync(
                artifacts.AssistantContextPath,
                screenshotId,
                $"_OCR failed: {exception.Message}_",
                CancellationToken.None);
        }
    }

    private async Task ReplaceOcrPlaceholderAsync(
        string assistantContextPath,
        string screenshotId,
        string replacement,
        CancellationToken cancellationToken)
    {
        var fileLock = contextFileLocks.GetOrAdd(
            NormalizeContextKey(assistantContextPath),
            _ => new SemaphoreSlim(1, 1));
        await fileLock.WaitAsync(cancellationToken);
        try
        {
            if (!File.Exists(assistantContextPath))
            {
                return;
            }

            var content = await File.ReadAllTextAsync(assistantContextPath, cancellationToken);
            var startMarker = $"<!-- screenshot-ocr:{screenshotId} -->";
            var endMarker = $"<!-- /screenshot-ocr:{screenshotId} -->";
            var startIndex = content.IndexOf(startMarker, StringComparison.Ordinal);
            var endIndex = content.IndexOf(endMarker, StringComparison.Ordinal);
            if (startIndex < 0 || endIndex < startIndex)
            {
                await File.AppendAllTextAsync(
                    assistantContextPath,
                    Environment.NewLine + replacement.TrimEnd() + Environment.NewLine,
                    cancellationToken);
                return;
            }

            endIndex += endMarker.Length;
            var updated = content[..startIndex] +
                replacement.TrimEnd() +
                content[endIndex..];
            await File.WriteAllTextAsync(assistantContextPath, updated, cancellationToken);
        }
        finally
        {
            fileLock.Release();
        }
    }

    private async Task<string> TryCreateCropMarkdownAsync(
        string assistantContextPath,
        string screenshotPath,
        ScreenshotCropCoordinates? crop,
        CancellationToken cancellationToken)
    {
        if (crop is null)
        {
            return "";
        }

        if (!OperatingSystem.IsWindows())
        {
            logger.LogWarning(
                "Ignoring screenshot crop coordinates for {ScreenshotPath} because image cropping is only supported on Windows",
                screenshotPath);
            return "";
        }

        try
        {
            var croppedPath = await SaveCroppedScreenshotAsync(screenshotPath, crop, cancellationToken);
            var relativePath = ToMarkdownPath(Path.GetRelativePath(
                Path.GetDirectoryName(assistantContextPath)!,
                croppedPath));
            return $"![Cropped screenshot]({relativePath})" +
                Environment.NewLine +
                Environment.NewLine;
        }
        catch (Exception exception) when (exception is InvalidDataException or ArgumentException or ExternalException)
        {
            logger.LogWarning(
                exception,
                "Ignoring invalid screenshot crop coordinates for {ScreenshotPath}",
                screenshotPath);
            return "";
        }
    }

#pragma warning disable CA1416
    private static Task<string> SaveCroppedScreenshotAsync(
        string screenshotPath,
        ScreenshotCropCoordinates crop,
        CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        using var original = new Bitmap(screenshotPath);
        ValidateCrop(crop, original.Width, original.Height);
        var cropRectangle = new Rectangle(crop.X, crop.Y, crop.Width, crop.Height);
        using var cropped = original.Clone(cropRectangle, original.PixelFormat);
        var croppedPath = Path.Combine(
            Path.GetDirectoryName(screenshotPath)!,
            $"{Path.GetFileNameWithoutExtension(screenshotPath)}-cropped.png");
        cropped.Save(croppedPath, ImageFormat.Png);
        return Task.FromResult(croppedPath);
    }
#pragma warning restore CA1416

    private static void ValidateCrop(ScreenshotCropCoordinates crop, int imageWidth, int imageHeight)
    {
        if (crop.X < 0 ||
            crop.Y < 0 ||
            crop.Width <= 0 ||
            crop.Height <= 0 ||
            crop.X + crop.Width > imageWidth ||
            crop.Y + crop.Height > imageHeight)
        {
            throw new InvalidDataException(
                $"Screenshot crop {crop.X},{crop.Y},{crop.Width},{crop.Height} is outside image bounds {imageWidth}x{imageHeight}.");
        }
    }

    private void TrackOcr(
        MeetingSessionArtifacts artifacts,
        CancellationTokenSource cancellation,
        Task task)
    {
        var contextKey = NormalizeContextKey(artifacts.AssistantContextPath);
        var tasks = pendingOcrByContext.GetOrAdd(contextKey, _ => []);
        lock (tasks)
        {
            tasks.RemoveAll(existing => existing.Task.IsCompleted);
            tasks.Add(new PendingOcrTask(task, cancellation));
        }
    }

    private static string CreateScreenshotMarkdown(
        string screenshotId,
        TimeSpan timestamp,
        string relativePath,
        string? windowTitle,
        bool ocrEnabled)
    {
        var title = string.IsNullOrWhiteSpace(windowTitle)
            ? ""
            : Environment.NewLine + $"Window: {windowTitle.Trim()}" + Environment.NewLine;
        var timestampText = FormatTimestamp(timestamp);
        var content = $"## Screenshot [{timestampText}]" +
            title +
            Environment.NewLine +
            $"![Screenshot {timestampText}]({relativePath})";
        if (!ocrEnabled)
        {
            return content;
        }

        return content +
            Environment.NewLine +
            Environment.NewLine +
            $"<!-- screenshot-ocr:{screenshotId} -->" +
            Environment.NewLine +
            "_OCR pending..._" +
            Environment.NewLine +
            $"<!-- /screenshot-ocr:{screenshotId} -->";
    }

    private static string ResolveAttachmentsFolder(string assistantContextPath, string configuredFolder)
    {
        var expanded = Environment.ExpandEnvironmentVariables(
            string.IsNullOrWhiteSpace(configuredFolder) ? "Attachments" : configuredFolder);
        return Path.IsPathRooted(expanded)
            ? Path.GetFullPath(expanded)
            : Path.GetFullPath(Path.Combine(Path.GetDirectoryName(assistantContextPath)!, expanded));
    }

    private static TimeSpan CalculateMeetingTimestamp(DateTimeOffset? meetingStartedAt, DateTimeOffset capturedAt)
    {
        if (meetingStartedAt is null)
        {
            return TimeSpan.Zero;
        }

        var elapsed = capturedAt - meetingStartedAt.Value;
        return elapsed < TimeSpan.Zero ? TimeSpan.Zero : elapsed;
    }

    private static string FormatTimestamp(TimeSpan timestamp)
    {
        return $"{(int)timestamp.TotalHours:00}:{timestamp.Minutes:00}:{timestamp.Seconds:00}";
    }

    private static string ToMarkdownPath(string path)
    {
        return path.Replace(Path.DirectorySeparatorChar, '/')
            .Replace(Path.AltDirectorySeparatorChar, '/');
    }

    private static string NormalizeContextKey(string assistantContextPath)
    {
        return Path.GetFullPath(assistantContextPath);
    }

    private sealed record PendingOcrTask(Task Task, CancellationTokenSource Cancellation);
}