meeting-assistant/MeetingAssistant/Screenshots/LiteLlmScreenshotOcrClient.cs

using System.Text.Json;
using System.Text.RegularExpressions;
using MeetingAssistant.MeetingNotes;
using MeetingAssistant.Summary;
using Microsoft.Extensions.AI;

namespace MeetingAssistant.Screenshots;

public sealed partial class LiteLlmScreenshotOcrClient : IScreenshotOcrClient
{
    private readonly ILogger<LiteLlmScreenshotOcrClient> logger;
    private readonly Func<HttpMessageHandler>? httpMessageHandlerFactory;

    public LiteLlmScreenshotOcrClient(ILogger<LiteLlmScreenshotOcrClient> logger)
    {
        this.logger = logger;
    }

    internal LiteLlmScreenshotOcrClient(
        Func<HttpMessageHandler> httpMessageHandlerFactory,
        ILogger<LiteLlmScreenshotOcrClient> logger)
    {
        this.httpMessageHandlerFactory = httpMessageHandlerFactory;
        this.logger = logger;
    }

    public async Task<ScreenshotOcrResult> ExtractAsync(
        string screenshotPath,
        string prompt,
        MeetingAssistantOptions options,
        CancellationToken cancellationToken)
    {
        var endpoint = !string.IsNullOrWhiteSpace(options.Screenshots.Ocr.Endpoint)
            ? options.Screenshots.Ocr.Endpoint
            : options.Agent.Endpoint;
        var model = !string.IsNullOrWhiteSpace(options.Screenshots.Ocr.Model)
            ? options.Screenshots.Ocr.Model
            : options.Agent.Model;
        var key = ResolveApiKey(options);
        var imageBytes = await File.ReadAllBytesAsync(screenshotPath, cancellationToken);
        var httpClient = CreateHttpClient();
        httpClient.BaseAddress = LiteLlmResponsesChatClient.NormalizeEndpoint(new Uri(endpoint));
        using var chatClient = new LiteLlmResponsesChatClient(
            httpClient,
            key,
            model,
            enableThinking: false,
            reasoningEffort: "none",
            reconnectionAttempts: options.Agent.ReconnectionAttempts,
            reconnectionDelay: options.Agent.ReconnectionDelay,
            logger: logger,
            firstRequestIsUser: false,
            useStreaming: options.Agent.UseStreaming);
        var response = await chatClient.GetResponseAsync(
            [
                new ChatMessage(
                    ChatRole.User,
                    [
                        new TextContent(CreatePrompt(prompt, imageBytes)),
                        new DataContent(imageBytes, "image/png")
                    ])
            ],
            cancellationToken: cancellationToken);
        var text = response.Text ?? string.Empty;
        logger.LogInformation("Screenshot OCR completed for {ScreenshotPath}", screenshotPath);
        return ParseOcrResult(text);
    }

    private HttpClient CreateHttpClient()
    {
        return httpMessageHandlerFactory is null
            ? new HttpClient()
            : new HttpClient(httpMessageHandlerFactory());
    }

    private static string CreatePrompt(string prompt, byte[] imageBytes)
    {
        return TryReadPngDimensions(imageBytes, out var width, out var height)
            ? prompt + Environment.NewLine + Environment.NewLine +
                $"Original screenshot dimensions: {width}x{height} pixels. Crop coordinates must fit within these bounds."
            : prompt;
    }

    private static ScreenshotOcrResult ParseOcrResult(string text)
    {
        ScreenshotCropCoordinates? crop = null;
        IReadOnlyList<string> attendees = [];
        var cleaned = JsonCodeBlockRegex().Replace(text, match =>
        {
            if (TryParseMetadata(match.Groups["json"].Value, out var parsedCrop, out var parsedAttendees))
            {
                crop = parsedCrop;
                attendees = parsedAttendees;
                return "";
            }

            return match.Value;
        }).Trim();
        return new ScreenshotOcrResult(cleaned, crop, attendees);
    }

    private static bool TryParseMetadata(
        string json,
        out ScreenshotCropCoordinates? crop,
        out IReadOnlyList<string> attendees)
    {
        crop = null;
        attendees = [];
        try
        {
            using var document = JsonDocument.Parse(json);
            if (document.RootElement.ValueKind != JsonValueKind.Object)
            {
                return false;
            }

            var hasMetadata = false;
            if (document.RootElement.TryGetProperty("attendees", out var attendeesElement))
            {
                hasMetadata = true;
                if (attendeesElement.ValueKind == JsonValueKind.Array)
                {
                    attendees = MeetingAttendeeNames.NormalizeDistinct(attendeesElement
                        .EnumerateArray()
                        .Where(element => element.ValueKind == JsonValueKind.String)
                        .Select(element => element.GetString() ?? ""));
                }
            }

            if (!document.RootElement.TryGetProperty("crop", out var cropElement))
            {
                return hasMetadata;
            }

            hasMetadata = true;
            if (cropElement.ValueKind == JsonValueKind.Null)
            {
                return true;
            }

            if (cropElement.ValueKind == JsonValueKind.Object &&
                TryGetInt(cropElement, "x", out var x) &&
                TryGetInt(cropElement, "y", out var y) &&
                TryGetInt(cropElement, "width", out var width) &&
                TryGetInt(cropElement, "height", out var height))
            {
                crop = new ScreenshotCropCoordinates(x, y, width, height);
            }

            return true;
        }
        catch (JsonException)
        {
            return false;
        }
    }

    private static bool TryGetInt(JsonElement element, string propertyName, out int value)
    {
        value = 0;
        return element.TryGetProperty(propertyName, out var property) &&
            property.ValueKind == JsonValueKind.Number &&
            property.TryGetInt32(out value);
    }

    private static bool TryReadPngDimensions(byte[] bytes, out int width, out int height)
    {
        width = 0;
        height = 0;
        if (bytes.Length < 24 ||
            bytes[0] != 0x89 ||
            bytes[1] != 0x50 ||
            bytes[2] != 0x4E ||
            bytes[3] != 0x47)
        {
            return false;
        }

        width = ReadBigEndianInt32(bytes, 16);
        height = ReadBigEndianInt32(bytes, 20);
        return width > 0 && height > 0;
    }

    private static int ReadBigEndianInt32(byte[] bytes, int offset)
    {
        return (bytes[offset] << 24) |
            (bytes[offset + 1] << 16) |
            (bytes[offset + 2] << 8) |
            bytes[offset + 3];
    }

    private static string ResolveApiKey(MeetingAssistantOptions options)
    {
        if (!string.IsNullOrWhiteSpace(options.Screenshots.Ocr.Key))
        {
            return options.Screenshots.Ocr.Key;
        }

        var ocrKey = Environment.GetEnvironmentVariable(options.Screenshots.Ocr.KeyEnv);
        if (!string.IsNullOrWhiteSpace(ocrKey))
        {
            return ocrKey;
        }

        if (!string.IsNullOrWhiteSpace(options.Agent.Key))
        {
            return options.Agent.Key;
        }

        var agentKey = Environment.GetEnvironmentVariable(options.Agent.KeyEnv);
        if (!string.IsNullOrWhiteSpace(agentKey))
        {
            return agentKey;
        }

        throw new InvalidOperationException(
            $"No screenshot OCR API key configured. Set MeetingAssistant:Screenshots:Ocr:Key or environment variable '{options.Screenshots.Ocr.KeyEnv}'.");
    }

    [GeneratedRegex("```json\\s*(?<json>.*?)\\s*```", RegexOptions.Singleline | RegexOptions.IgnoreCase)]
    private static partial Regex JsonCodeBlockRegex();
}