SkiaStitch.cs

using AForge;
using BioLib;
using BruTile;
using CSScripting;
using OpenSlideGTK;
using org.apache.xmlgraphics.image.codec.tiff;
using SkiaSharp;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
 
namespace BioGTK
{
    public class SkiaStitchingPipeline : IDisposable
    {
        private const bool VerboseLogging = false;
        #region Fields and Properties
 
        private readonly SlideBase _slideSource;
        private readonly OpenSlideGTK.OpenSlideBase _openSlideSource;
        private readonly bool _isOpenSlide;
 
        // Tile cache with thread-safe access
        private ConcurrentDictionary<string, CachedTile> _tileCache;
        private int _maxCacheSize;
        private object _cacheLock = new object();
 
        // Rendering state
        private SKSurface _compositeSurface;
        private SKCanvas _compositeCanvas;
        private int _currentSurfaceWidth = 600;
        private int _currentSurfaceHeight = 400;
 
        // Async operation management
        private ConcurrentDictionary<string, Task<SKImage>> _pendingTileFetches;
 
        // Configuration
        public int TileSize { get; set; } = 256;
        public int PrefetchRadius { get; set; } = 1;
        public SKSamplingOptions Sampling { get; set; } = SKSamplingOptions.Default;
        public bool EnablePrefetch { get; set; } = false;
 
        #endregion
 
        #region Nested Types
 
        private class CachedTile
        {
            public SKImage Image { get; set; }
            public DateTime LastAccessed { get; set; }
            public long SizeBytes { get; set; }
            public string Key { get; set; }
 
            public CachedTile(SKImage image, string key)
            {
                Image = image;
                Key = key;
                LastAccessed = DateTime.UtcNow;
                SizeBytes = CalculateSize(image);
            }
 
            private long CalculateSize(SKImage image)
            {
                if (image == null) return 0;
                return image.Width * image.Height * 4; // Assume RGBA
            }
        }
 
        public class TileRequest
        {
            public int Level
            {
                get => Index.Level;
                set => Index = new TileIndex(Index.Col, Index.Row, value);
            }
 
            // FIXED: Col is X, Row is Y
            public int TileX
            {
                get => Index.Col * 256;  // Col = X coordinate
                set
                {
                    int col = value / 256;
                    Index = new TileIndex(col, Index.Row, Level);
                }
            }
 
            public int TileY
            {
                get => Index.Row * 256;  // Row = Y coordinate
                set
                {
                    int row = value / 256;
                    Index = new TileIndex(Index.Col, row, Level);
                }
            }
 
            public ZCT Coordinate { get; set; }
            public Extent Extent { get; set; }
            public TileIndex Index { get; set; }
 
            public RectangleD Bounds
            {
                get
                {
                    return new RectangleD(
                        Extent.MinX,
                        Extent.MinY,
                        Extent.MaxX - Extent.MinX,
                        Extent.MaxY - Extent.MinY
                    );
                }
            }
 
            public TileInfo TileInfo
            {
                get => new TileInfo { Index = Index, Extent = Extent };
                set
                {
                    Index = value.Index;
                    Extent = value.Extent;
                }
            }
 
            public TileRequest(TileInfo tf, ZCT coord)
            {
                Index = tf.Index;
                Extent = tf.Extent;
                Coordinate = coord;
            }
 
            public string GetCacheKey()
            {
                // Use Col/Row directly to avoid confusion
                string rgbKey = string.Empty;
                if (App.viewer?.Mode == ImageView.ViewMode.RGBImage && ImageView.SelectedImage != null)
                {
                    rgbKey = $"_{ImageView.SelectedImage.rgbChannels[0]}_{ImageView.SelectedImage.rgbChannels[1]}_{ImageView.SelectedImage.rgbChannels[2]}";
                }
                return $"{Index.Col}_{Index.Row}_{Level}_{Coordinate.Z}_{Coordinate.C}_{Coordinate.T}{rgbKey}";
            }
        }
        #endregion
 
        #region Constructor and Initialization
 
        public SkiaStitchingPipeline(OpenSlideGTK.OpenSlideBase openSlideSource, int maxCacheSizeMB = 512)
        {
            _openSlideSource = openSlideSource ?? throw new ArgumentNullException(nameof(openSlideSource));
            _isOpenSlide = true;
 
            InitializeCommon(maxCacheSizeMB);
        }
 
        public SkiaStitchingPipeline(SlideBase slideSource, int maxCacheSizeMB = 512)
        {
            _slideSource = slideSource ?? throw new ArgumentNullException(nameof(slideSource));
            _isOpenSlide = false;
 
            InitializeCommon(maxCacheSizeMB);
        }
 
        private void InitializeCommon(int maxCacheSizeMB)
        {
            _maxCacheSize = maxCacheSizeMB * 1024 * 1024;
            _tileCache = new ConcurrentDictionary<string, CachedTile>();
            _pendingTileFetches = new ConcurrentDictionary<string, Task<SKImage>>();
        }
 
        #endregion
 
        #region Public Stitching Methods
        private readonly SemaphoreSlim _renderSemaphore = new SemaphoreSlim(1, 1);
        public async Task<SKImage> StitchViewportAsync(
            PointD origin,
            double resolution,
            ZCT coordinate,
            int viewportWidth = 600,
            int viewportHeight = 400)
        {
            // Calculate viewport parameters
            int level = CalculateLevelFromResolution(resolution);
            var tileRequests = CalculateTileRequests(origin, viewportWidth, viewportHeight, level, coordinate, resolution);
            if (viewportWidth <= 1 || viewportHeight <= 1)
            {
                viewportWidth = 600;
                viewportHeight = 400;
            }
            // Create a NEW surface for this render - no shared state, fully parallel-safe
            using (var surface = SKSurface.Create(new SKImageInfo(
                viewportWidth,
                viewportHeight,
                SKColorType.Rgba8888,
                SKAlphaType.Premul)))
            {
                if (surface == null)
                {
                    Console.WriteLine("Failed to create SKSurface");
                    return null;
                }
 
                var canvas = surface.Canvas;
                canvas.Clear(SKColors.Gray);
 
                // Fetch and composite all tiles
                await CompositeTilesAsync(canvas, tileRequests, origin, resolution);
 
                // Create and return snapshot
                return surface.Snapshot();
            }  // Surface automatically disposed here
        }
 
        public SKImage StitchRegion(
            RectangleD region,
            int level,
            ZCT coordinate, double resolution, int pxwidth, int pxheight)
        {
            var tileRequests = CalculateTileRequests(
                new PointD(region.X,region.Y),pxwidth, pxheight, level, coordinate, resolution);
 
            int outputWidth = (int)Math.Ceiling(region.W);
            int outputHeight = (int)Math.Ceiling(region.H);
 
            using (var surface = SKSurface.Create(new SKImageInfo(outputWidth, outputHeight)))
            {
                var canvas = surface.Canvas;
                canvas.Clear(SKColors.Transparent);
 
                foreach (var request in tileRequests)
                {
                    var tileImage = FetchTileSync(request);
                    if (tileImage != null)
                    {
                        DrawTileToCanvas(canvas,request, new AForge.PointD(region.X,region.Y), resolution,  tileImage);
                    }
                }
 
                return surface.Snapshot();
            }
        }
 
        #endregion
 
        #region Tile Fetching and Caching
 
        private async Task<SKImage> FetchTileAsync(TileRequest request, ZCT coordinate)
        {
            string key = request.GetCacheKey();
 
            // 1. Check cache again to be sure
            if (_tileCache.TryGetValue(key, out var cached)) return cached.Image;
 
            // 2. Check if already being fetched
            if (_pendingTileFetches.TryGetValue(key, out var pendingTask)) return await pendingTask;
 
            // 3. Create the fetch task
            var fetchTask = Task.Run(async () =>
            {
                try
                {
                    byte[] data;
                    if (_isOpenSlide)
                    {
                        TileInfo tf = new TileInfo { Index = request.Index, Extent = request.Extent };
                        // OpenSlide usually takes Level, X-index, Y-index
                        data = _openSlideSource.GetTile(tf);
                    }
                    else
                    {
                        TileInfo tf = new TileInfo { Index = request.Index, Extent = request.Extent };
                        data = await FetchSlideTileAsync(tf, request.Coordinate).ConfigureAwait(false);
                    }
 
                    if (data == null || data.Length == 0) return null;
 
                    // Convert to Skia Image (using the shear-fix logic from before)
                    var levelResolution = CalculateResolutionFromLevel(request.Index.Level);
                    var image = ConvertTileDataToSKImage(data, request);
 
                    if (image != null)
                    {
                        var cachedTile = new CachedTile(image, DateTime.Now.ToString());
                        _tileCache.AddOrUpdate(key, cachedTile, (k, old) => cachedTile);
                        return image;
                    }
                }
                catch (Exception ex)
                {
                    if (VerboseLogging) Console.WriteLine($"Fetch error for tile {key}: {ex.Message} + {ex.StackTrace}");
                }
                finally
                {
                    _pendingTileFetches.TryRemove(key, out _);
                }
                return null;
            });
 
            _pendingTileFetches[key] = fetchTask;
            return await fetchTask;
        }
 
        private bool ShouldComposeRgbTile()
        {
            if (_isOpenSlide || App.viewer?.Mode != ImageView.ViewMode.RGBImage)
                return false;
 
            var bioImage = _slideSource?.Image?.BioImage;
            if (bioImage == null || bioImage.Channels == null || bioImage.Channels.Count <= 1)
                return false;
            if (bioImage.Resolutions == null || bioImage.Resolutions.Count == 0)
                return false;
 
            var pixelFormat = bioImage.Resolutions[0].PixelFormat;
            return pixelFormat != AForge.PixelFormat.Format24bppRgb &&
                   pixelFormat != AForge.PixelFormat.Format48bppRgb;
        }
 
        private async Task<byte[]> FetchSlideTileAsync(TileInfo tileInfo, ZCT coordinate)
        {
            if (!ShouldComposeRgbTile())
                return _slideSource.GetTile(tileInfo, coordinate);
 
            var bioImage = _slideSource?.Image?.BioImage;
            if (bioImage == null)
                return _slideSource.GetTile(tileInfo, coordinate);
 
            int channelCount = Math.Max(1, bioImage.Channels.Count);
            int redIndex = Math.Clamp(bioImage.rgbChannels[0], 0, channelCount - 1);
            int greenIndex = Math.Clamp(bioImage.rgbChannels[1], 0, channelCount - 1);
            int blueIndex = Math.Clamp(bioImage.rgbChannels[2], 0, channelCount - 1);
 
            byte[] redTile = await _slideSource.GetTileAsync(new BioLib.TileInformation(tileInfo.Index, tileInfo.Extent, new ZCT(coordinate.Z, redIndex, coordinate.T)), new ZCT(coordinate.Z, redIndex, coordinate.T)).ConfigureAwait(false);
            byte[] greenTile = await _slideSource.GetTileAsync(new BioLib.TileInformation(tileInfo.Index, tileInfo.Extent, new ZCT(coordinate.Z, greenIndex, coordinate.T)), new ZCT(coordinate.Z, greenIndex, coordinate.T)).ConfigureAwait(false);
            byte[] blueTile = await _slideSource.GetTileAsync(new BioLib.TileInformation(tileInfo.Index, tileInfo.Extent, new ZCT(coordinate.Z, blueIndex, coordinate.T)), new ZCT(coordinate.Z, blueIndex, coordinate.T)).ConfigureAwait(false);
            return ComposeBgraTiles(redTile, greenTile, blueTile);
        }
 
        private static byte[] ComposeBgraTiles(byte[] redTile, byte[] greenTile, byte[] blueTile)
        {
            byte[] fallback = redTile ?? greenTile ?? blueTile;
            if (fallback == null)
                return null;
 
            redTile ??= fallback;
            greenTile ??= fallback;
            blueTile ??= fallback;
 
            int byteCount = Math.Min(redTile.Length, Math.Min(greenTile.Length, blueTile.Length));
            byteCount -= byteCount % 4;
            if (byteCount <= 0)
                return null;
 
            byte[] composed = new byte[byteCount];
            for (int i = 0; i < byteCount; i += 4)
            {
                byte r = ExtractIntensity(redTile, i);
                byte g = ExtractIntensity(greenTile, i);
                byte b = ExtractIntensity(blueTile, i);
                byte a = (byte)Math.Max(redTile[i + 3], Math.Max(greenTile[i + 3], blueTile[i + 3]));
                if (a == 0 && (r != 0 || g != 0 || b != 0))
                    a = byte.MaxValue;
 
                composed[i] = b;
                composed[i + 1] = g;
                composed[i + 2] = r;
                composed[i + 3] = a;
            }
 
            return composed;
        }
 
        private static byte ExtractIntensity(byte[] tile, int offset)
        {
            return (byte)Math.Max(tile[offset], Math.Max(tile[offset + 1], tile[offset + 2]));
        }
 
        private SKImage FetchTileSync(TileRequest request)
        {
            string cacheKey = request.GetCacheKey();
 
            // Check cache
            if (_tileCache.TryGetValue(cacheKey, out var cachedTile))
            {
                cachedTile.LastAccessed = DateTime.UtcNow;
                return cachedTile.Image;
            }
 
            // Fetch directly
            try
            {
                byte[] tileData = FetchTileDataSync(request);
                if (tileData == null || tileData.Length == 0)
                    return null;
 
                SKImage skImage = ConvertTileDataToSKImage(tileData, request);
 
                if (skImage != null)
                {
                    AddToCache(cacheKey, skImage);
                }
 
                return skImage;
            }
            catch (Exception ex)
            {
                if (VerboseLogging) Console.WriteLine($"Error fetching tile sync {cacheKey}: {ex.Message}");
                return null;
            }
        }
 
        private async Task<byte[]> FetchTileDataAsync(TileRequest request)
        {
            if (VerboseLogging)
            {
                Console.WriteLine($"\n--- Fetching tile data for Level {request.Level}, Index[{request.Index.Col},{request.Index.Row}] ---");
                Console.WriteLine($"  Extent: ({request.Extent.MinX:F0}, {request.Extent.MinY:F0}) to ({request.Extent.MaxX:F0}, {request.Extent.MaxY:F0})");
            }
 
            try
            {
                if (_isOpenSlide)
                {
                    // Use BruTile's GetTile method directly
                    byte[] tileData = _openSlideSource.GetTile(request.TileInfo);
 
                    if (tileData == null || tileData.Length == 0)
                    {
                        if (VerboseLogging) Console.WriteLine($"  ERROR: OpenSlide returned null/empty tile data");
                        return null;
                    }
 
                    if (VerboseLogging) Console.WriteLine($"  OpenSlide returned {tileData.Length} bytes");
                    return tileData;
                }
                else
                {
                    // For BioLib SlideBase
                    var sliceInfo = new BioLib.SliceInfo(
                        request.Extent.MinX,
                        request.Extent.MinY,
                        request.Bounds.W,
                        request.Bounds.H,
                        CalculateResolutionFromLevel(request.Level),
                        request.Coordinate
                    );
 
                    var origin = new PointD(request.Extent.MinX, request.Extent.MinY);
                    var size = new AForge.Size((int)request.Bounds.W, (int)request.Bounds.H);
                    TileInformation tf = new TileInformation(request.Index, request.Extent, request.Coordinate);
                    byte[] tileData = await _slideSource.GetTileAsync(tf, request.Coordinate);
 
                    if (tileData == null || tileData.Length == 0)
                    {
                        if (VerboseLogging) Console.WriteLine($"  ERROR: SlideBase returned null/empty tile data");
                        return null;
                    }
 
                    if (VerboseLogging) Console.WriteLine($"  SlideBase returned {tileData.Length} bytes");
                    return tileData;
                }
            }
            catch (Exception ex)
            {
                if (VerboseLogging)
                {
                    Console.WriteLine($"  EXCEPTION in FetchTileDataAsync: {ex.Message}");
                    Console.WriteLine($"  {ex.StackTrace}");
                }
                return null;
            }
        }
        private byte[] FetchTileDataSync(TileRequest request)
        {
            if (_isOpenSlide)
            {
                var sliceInfo = new OpenSlideGTK.SliceInfo(
                    request.Bounds.X,
                    request.Bounds.Y,
                    request.Bounds.W,
                    request.Bounds.H,
                    CalculateResolutionFromLevel(request.Level)
                );
 
                return _openSlideSource.GetSlice(sliceInfo, request.Coordinate, request.Level);
            }
            else
            {
                if (ShouldComposeRgbTile())
                {
                    TileInfo tf = new TileInfo { Index = request.Index, Extent = request.Extent };
                    return FetchSlideTileAsync(tf, request.Coordinate).Result;
                }
 
                var sliceInfo = new BioLib.SliceInfo(
                    request.Bounds.X,
                    request.Bounds.Y,
                    request.Bounds.W,
                    request.Bounds.H,
                    CalculateResolutionFromLevel(request.Level),
                    request.Coordinate
                );
 
                var origin = new PointD(request.Bounds.X, request.Bounds.Y);
                var size = new AForge.Size((int)request.Bounds.W, (int)request.Bounds.H);
 
                return _slideSource.GetSlice(sliceInfo, origin, size).Result;
            }
        }
 
        #endregion
 
        #region Tile Compositing
 
        private async Task CompositeTilesAsync(
            SKCanvas canvas,
            List<TileRequest> tileRequests,
            PointD origin,
            double resolution)
        {
            if (tileRequests.Count == 0) return;
 
            // Fetch all tiles in parallel
            var tileTasks = tileRequests.Select(async req =>
            {
                var image = await FetchTileAsync(req, req.Coordinate);
                return new { Request = req, Image = image };
            }).ToArray();
 
            var results = await Task.WhenAll(tileTasks);
 
            using (var paint = new SKPaint { IsAntialias = false })
            {
                foreach (var result in results)
                {
                    if (result.Image == null) continue;
                    DrawTileToCanvas(canvas, result.Request, origin, resolution, result.Image, paint);
                }
            }
        }
        private void DrawTileToCanvas(
            SKCanvas canvas,
            TileRequest request,
            PointD viewportOrigin,
            double resolution,
            SKImage tileImage,
            SKPaint paint = null)
        {
            if (tileImage == null) return;
 
            // X: (Tile Left - Viewport Left) / Resolution
            float destX = (float)((request.Extent.MinX - viewportOrigin.X) / resolution);
 
            // Viewport origin is stored in positive-down image/stage space.
            // BruTile extents remain in OSM world space, so convert tile top
            // into the same positive-down convention here.
            float destY = (float)((-request.Extent.MaxY - viewportOrigin.Y) / resolution);
 
            // Calculate dimensions based on the Extent to ensure precision
            float destW = (float)(request.Extent.Width / resolution);
            float destH = (float)(request.Extent.Height / resolution);
 
            var destRect = new SKRect(destX, destY, destX + destW, destY + destH);
 
            // Use Sampling if provided in the class, otherwise use the paint object
            if (paint != null)
                canvas.DrawImage(tileImage, destRect, paint);
            else
                canvas.DrawImage(tileImage, destRect, Sampling);
        }
        private SKPaint CreateCompositePaint()
        {
            return new SKPaint
            {
                IsAntialias = true,
                BlendMode = SKBlendMode.SrcOver
            };
        }
 
        private SKPaint CreateTilePaint()
        {
            return new SKPaint
            {
                IsAntialias = false // Tiles should align perfectly
            };
        }
 
        #endregion
 
        #region Tile Calculation and Management
 
        private List<TileRequest> CalculateTileRequests(PointD origin, int viewportWidth, int viewportHeight, int level, ZCT coordinate, double resolution)
        {
            // origin is the Top-Left of the viewport in positive-down image/stage
            // space. Convert to BruTile's OSM world-space extent here.
 
            double minX = origin.X;
            double maxX = origin.X + (viewportWidth * resolution);
            double maxY = -origin.Y;
            double minY = -origin.Y - (viewportHeight * resolution);
 
            // Create the extent (MinX, MinY, MaxX, MaxY)
            var viewportExtent = new Extent(minX, minY, maxX, maxY);
 
            IEnumerable<TileInfo> tileInfos;
            if (_isOpenSlide)
            {
                tileInfos = _openSlideSource.Schema.GetTileInfos(viewportExtent, level);
            }
            else
            {
                tileInfos = _slideSource.Schema.GetTileInfos(viewportExtent, level);
            }
 
            var requests = new List<TileRequest>();
            foreach (var info in tileInfos)
            {
                requests.Add(new TileRequest(info, coordinate));
            }
            return requests;
        }
        #endregion
 
        #region Cache Management
 
        private void AddToCache(string key, SKImage image)
        {
            var cachedTile = new CachedTile(image, key);
 
            lock (_cacheLock)
            {
                _tileCache[key] = cachedTile;
 
                // Check if we need to evict
                long currentCacheSize = CalculateTotalCacheSize();
 
                if (currentCacheSize > _maxCacheSize)
                {
                    EvictLeastRecentlyUsedTiles(currentCacheSize - _maxCacheSize);
                }
            }
        }
 
        private void EvictLeastRecentlyUsedTiles(long bytesToFree)
        {
            var sortedTiles = new List<CachedTile>(_tileCache.Values);
            sortedTiles.Sort((a, b) => a.LastAccessed.CompareTo(b.LastAccessed));
 
            long freedBytes = 0;
            foreach (var tile in sortedTiles)
            {
                if (freedBytes >= bytesToFree)
                    break;
 
                if (_tileCache.TryRemove(tile.Key, out var removed))
                {
                    freedBytes += removed.SizeBytes;
                    removed.Image?.Dispose();
                }
            }
        }
 
        private long CalculateTotalCacheSize()
        {
            long total = 0;
            foreach (var tile in _tileCache.Values)
            {
                total += tile.SizeBytes;
            }
            return total;
        }
 
        public void ClearCache()
        {
            lock (_cacheLock)
            {
                foreach (var tile in _tileCache.Values)
                {
                    tile.Image?.Dispose();
                }
                _tileCache.Clear();
            }
        }
 
        #endregion
 
        #region Utility Methods
 
        // SkiaStitch.cs - Update this method
        private SKImage ConvertTileDataToSKImage(byte[] tileData, TileRequest request)
        {
            try
            {
                // 1. Try standard decoding (JPEG/PNG/etc)
                using (var data = SKData.CreateCopy(tileData))
                {
                    var image = SKImage.FromEncodedData(data);
                    if (image != null) return image;
                }
 
                // 2. Fallback for Raw Pixels (The "Shear" Fix)
                // Use the Schema's defined TileWidth/Height, NOT the calculated resolution width
                int width = _isOpenSlide ? _openSlideSource.Schema.GetTileWidth(request.Index.Level)
                                         : _slideSource.Schema.GetTileWidth(request.Index.Level);
                int height = _isOpenSlide ? _openSlideSource.Schema.GetTileHeight(request.Index.Level)
                                          : _slideSource.Schema.GetTileHeight(request.Index.Level);
 
                var info = new SKImageInfo(width, height, SKColorType.Bgra8888, SKAlphaType.Premul);
                // If decode fails, create from raw RGBA data
                // CRITICAL: Don't dispose SKData until image is disposed
                var imageInfo = new SKImageInfo(width, height, SKColorType.Rgba8888, SKAlphaType.Premul);
                var rawData = SKData.CreateCopy(tileData);
                var pixmap = new SKPixmap(imageInfo, rawData.Data, imageInfo.RowBytes);
 
                // Use release delegate to properly manage lifetime
                return SKImage.FromPixels(pixmap, (address, context) =>
                {
                    rawData?.Dispose();
                }, null);
            }
            catch (Exception ex)
            {
                if (VerboseLogging) Console.WriteLine($"Error converting tile data: {ex.Message}");
                return null;
            }
        }
        private int CalculateLevelFromResolution(double resolution)
        {
            if (_isOpenSlide)
            {
                return OpenSlideGTK.TileUtil.GetLevel(_openSlideSource.Schema.Resolutions, resolution);
            }
            else
            {
                return OpenSlideGTK.TileUtil.GetLevel(_slideSource.Schema.Resolutions, resolution);
            }
        }
 
        private double CalculateResolutionFromLevel(int level)
        {
            if (_isOpenSlide)
            {
                return _openSlideSource.Schema.Resolutions[level].UnitsPerPixel;
            }
            else
            {
                return _slideSource.Schema.Resolutions[level].UnitsPerPixel;
            }
        }
 
        #endregion
 
        #region Disposal
        public void Dispose()
        {
            ClearCache();
            _pendingTileFetches.Clear();;
        }
        #endregion
    }
}