meeting-assistant/MeetingAssistant/Screenshots/LiteLlmScreenshotOcrClient.cs

using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using System.Text.RegularExpressions;
using MeetingAssistant.MeetingNotes;

namespace MeetingAssistant.Screenshots;

public sealed partial class LiteLlmScreenshotOcrClient : IScreenshotOcrClient
{
    private static readonly JsonSerializerOptions JsonOptions = new(JsonSerializerDefaults.Web);
    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);
        using var httpClient = CreateHttpClient();
        httpClient.BaseAddress = NormalizeEndpoint(new Uri(endpoint));
        httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", key);
        var payload = CreatePayload(model, CreatePrompt(prompt, imageBytes), imageBytes);
        using var content = new StringContent(payload.ToJsonString(JsonOptions), Encoding.UTF8, "application/json");
        using var response = await httpClient.PostAsync("responses", content, cancellationToken);
        var responseJson = await response.Content.ReadAsStringAsync(cancellationToken);
        if (!response.IsSuccessStatusCode)
        {
            throw new InvalidOperationException(
                $"Screenshot OCR request failed with {(int)response.StatusCode} {response.ReasonPhrase}: {responseJson}");
        }

        var text = ParseOutputText(responseJson);
        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 JsonObject CreatePayload(string model, string prompt, byte[] imageBytes)
    {
        return new JsonObject
        {
            ["model"] = model,
            ["store"] = false,
            ["input"] = new JsonArray
            {
                new JsonObject
                {
                    ["role"] = "user",
                    ["content"] = new JsonArray
                    {
                        new JsonObject
                        {
                            ["type"] = "input_text",
                            ["text"] = prompt
                        },
                        new JsonObject
                        {
                            ["type"] = "input_image",
                            ["image_url"] = "data:image/png;base64," + Convert.ToBase64String(imageBytes)
                        }
                    }
                }
            }
        };
    }

    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 ParseOutputText(string responseJson)
    {
        using var document = JsonDocument.Parse(responseJson);
        var root = document.RootElement;
        var parts = new List<string>();
        if (root.TryGetProperty("output_text", out var outputText) &&
            outputText.ValueKind == JsonValueKind.String &&
            !string.IsNullOrWhiteSpace(outputText.GetString()))
        {
            parts.Add(outputText.GetString()!);
        }

        if (root.TryGetProperty("output", out var output) && output.ValueKind == JsonValueKind.Array)
        {
            foreach (var item in output.EnumerateArray())
            {
                if (!item.TryGetProperty("content", out var content) || content.ValueKind != JsonValueKind.Array)
                {
                    continue;
                }

                foreach (var block in content.EnumerateArray())
                {
                    if (block.TryGetProperty("type", out var type) &&
                        type.GetString() == "output_text" &&
                        block.TryGetProperty("text", out var text) &&
                        text.ValueKind == JsonValueKind.String &&
                        !string.IsNullOrWhiteSpace(text.GetString()))
                    {
                        parts.Add(text.GetString()!);
                    }
                }
            }
        }

        return parts.Count == 0
            ? ""
            : string.Join(Environment.NewLine + Environment.NewLine, parts);
    }

    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}'.");
    }

    private static Uri NormalizeEndpoint(Uri endpoint)
    {
        var value = endpoint.ToString().TrimEnd('/');
        if (!value.EndsWith("/v1", StringComparison.OrdinalIgnoreCase))
        {
            value += "/v1";
        }

        return new Uri(value + "/");
    }

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