// The DirectX capture code is based off examples from the
// SharpDX Samples at https://github.com/sharpdx/SharpDX.

// Copyright (c) 2010-2013 SharpDX - Alexandre Mutel
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

using Insight.Remote.Shared.Abstractions;
using Insight.Remote.Shared.Extensions;
using Insight.Remote.Shared.Models;
using Insight.Remote.Shared.Native.Windows;
using Insight.Remote.Windows.Models;
using Microsoft.Extensions.Logging;
using Microsoft.Win32;
using RemoteControl.Shared;
using SharpDX;
using SharpDX.Direct3D11;
using SharpDX.DXGI;
using SkiaSharp;
using SkiaSharp.Views.Desktop;
using System.Diagnostics.CodeAnalysis;
using System.Drawing.Imaging;
using System.Runtime.InteropServices;
using Vaitr.Bus;
using Result = RemoteControl.Shared.Result;

namespace Insight.Remote.Windows.Services;

internal class WinScreenCapturer : IScreenProvider
{
    public event EventHandler<Rectangle>? ScreenChanged;
    public Rectangle GetVirtualScreenBounds() => SystemInformation.VirtualScreen;

    private SKBitmap? CurrentFrame
    {
        get => _currentFrame;
        set
        {
            if (_currentFrame != null)
            {
                _previousFrame?.Dispose();
                _previousFrame = _currentFrame;
            }
            _currentFrame = value;
        }
    }

    public int GetScreenCount() => Screen.AllScreens.Length;
    public IEnumerable<string> GetDisplayNames() => Screen.AllScreens.Select(x => x.DeviceName);
    public Rectangle CurrentScreenBounds { get; private set; } = Screen.PrimaryScreen?.Bounds ?? Rectangle.Empty;
    public string SelectedScreen { get; private set; } = Screen.PrimaryScreen?.DeviceName ?? string.Empty;

    private readonly object _screenBoundsLock = new();
    private SKBitmap? _currentFrame;
    private SKBitmap? _previousFrame;
    private bool _initialized;

    private readonly SemaphoreSlim _semaphore = new(1, 1);
    private readonly Dictionary<string, int> _bitBltScreens = new();
    private readonly Dictionary<string, DirectXOutput> _directxScreens = new();

    private readonly Bus _bus;
    private readonly ILogger<WinScreenCapturer> _logger;

    public WinScreenCapturer(Bus bus, ILogger<WinScreenCapturer> logger)
    {
        _bus = bus;
        _logger = logger;

        SystemEvents.DisplaySettingsChanged += SystemEvents_DisplaySettingsChanged;
    }

    public async Task InitAsync(CancellationToken cancellationToken)
    {
        Win32Interop.SwitchToInputDesktop();

        InitBitBlt();
        InitDirectX();

        _initialized = true;
    }

    public void SetSelectedScreen(string displayName)
    {
        lock (_screenBoundsLock)
        {
            if (displayName == SelectedScreen) return;

            if (_bitBltScreens.ContainsKey(displayName)) SelectedScreen = displayName;
            else SelectedScreen = _bitBltScreens.Keys.First();

            RefreshCurrentScreenBounds();
        }
    }

    public async Task<FrameResult> GetNextFrameAsync(CancellationToken cancellationToken)
    {
        try
        {
            await _semaphore.WaitAsync(cancellationToken);

            if (Win32Interop.SwitchToInputDesktop() is false)
            {
                // Something will occasionally prevent this from succeeding after active
                // desktop has changed to/from WinLogon (err code 170).  I'm guessing a hook
                // is getting put in the desktop, which causes SetThreadDesktop to fail.
                // The caller can start a new thread, which seems to resolve it.
                var errCode = Marshal.GetLastWin32Error();

                _logger.LogError("Failed to switch to input desktop. Last Win32 error code: {errCode}", errCode);

                return new FrameResult(null, false, true);
            }

            if (_initialized is false)
            {
                _logger.LogWarning("Init needed in GetNextFrame.");
                await InitAsync(cancellationToken);
            }

            var frame = GetDirectXFrame();
            if (frame is null || frame.IsSuccess is false)
            {
                var bitBltResult = GetBitBltFrame();
                if (bitBltResult.IsSuccess is false)
                {
                    var ex = bitBltResult.Exception ?? new("Unknown error.");
                    _logger.LogError(ex, "Error while getting next frame.");
                    return new FrameResult(null, false, true);
                }

                CurrentFrame = bitBltResult.Value;
            }
            else if (frame.IsSuccess && frame.Bitmap is not null && IsEmpty(frame.Bitmap) is false)
            {
                CurrentFrame = frame.Bitmap;
            }

            return new FrameResult(CurrentFrame, frame is not null && frame.HadChanges, false);
        }
        catch (TaskCanceledException)
        {
            return new FrameResult(null, false, true);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error while getting next frame.");
            _initialized = false;

            return new FrameResult(null, false, true);
        }
        finally
        {
            _semaphore.Release();
        }
    }

    public SKRect GetFrameDiffArea(bool fullscreen)
    {
        if (CurrentFrame is null) return SKRect.Empty;
        return CurrentFrame.GetDiffArea(_previousFrame, fullscreen);
    }

    public async Task<SKRect> GetDiffAreaAsync(bool fullscreen, CancellationToken cancellationToken)
    {
        if (CurrentFrame is null) return SKRect.Empty;

        try
        {
            await _semaphore.WaitAsync(cancellationToken);

            return CurrentFrame.GetDiffArea(_previousFrame, fullscreen);
        }
        finally
        {
            _semaphore.Release();
        }
    }

    public int GetSelectedScreenIndex()
    {
        if (_bitBltScreens.TryGetValue(SelectedScreen, out var index)) return index;
        return 0;
    }

    private Result<SKBitmap> GetBitBltFrame()
    {
        try
        {
            using var bitmap = new Bitmap(CurrentScreenBounds.Width, CurrentScreenBounds.Height, PixelFormat.Format32bppArgb);
            using (var graphic = Graphics.FromImage(bitmap))
            {
                graphic.CopyFromScreen(CurrentScreenBounds.Left, CurrentScreenBounds.Top, 0, 0, new Size(CurrentScreenBounds.Width, CurrentScreenBounds.Height));
            }
            return Result.Ok(bitmap.ToSKBitmap());
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Capturer error in BitBltCapture.");
            _initialized = false;
            return Result.Fail<SKBitmap>("Error while capturing BitBlt frame.");
        }
    }

    private DxCaptureResult GetDirectXFrame()
    {
        if (_directxScreens.TryGetValue(SelectedScreen, out var dxOutput) is false)
        {
            return DxCaptureResult.Fail("DirectX output not found.");
        }

        try
        {
            var outputDuplication = dxOutput.OutputDuplication;
            var device = dxOutput.Device;
            var texture2D = dxOutput.Texture2D;
            var bounds = dxOutput.Bounds;

            var result = outputDuplication.TryAcquireNextFrame(timeoutInMilliseconds: 25, out var duplicateFrameInfo, out var screenResource);

            if (!result.Success)
            {
                return DxCaptureResult.TryAcquireFailed(result);
            }

            if (duplicateFrameInfo.AccumulatedFrames == 0)
            {
                try
                {
                    outputDuplication.ReleaseFrame();
                }
                catch { }
                return DxCaptureResult.NoAccumulatedFrames(result);
            }

            using Texture2D screenTexture2D = screenResource.QueryInterface<Texture2D>();
            device.ImmediateContext.CopyResource(screenTexture2D, texture2D);
            var dataBox = device.ImmediateContext.MapSubresource(texture2D, 0, MapMode.Read, SharpDX.Direct3D11.MapFlags.None);
            using var bitmap = new Bitmap(bounds.Width, bounds.Height, PixelFormat.Format32bppArgb);
            var bitmapData = bitmap.LockBits(bounds, ImageLockMode.WriteOnly, bitmap.PixelFormat);
            var dataBoxPointer = dataBox.DataPointer;
            var bitmapDataPointer = bitmapData.Scan0;
            for (var y = 0; y < bounds.Height; y++)
            {
                Utilities.CopyMemory(bitmapDataPointer, dataBoxPointer, bounds.Width * 4);
                dataBoxPointer = IntPtr.Add(dataBoxPointer, dataBox.RowPitch);
                bitmapDataPointer = IntPtr.Add(bitmapDataPointer, bitmapData.Stride);
            }
            bitmap.UnlockBits(bitmapData);
            device.ImmediateContext.UnmapSubresource(texture2D, 0);
            screenResource?.Dispose();

            switch (dxOutput.Rotation)
            {
                case DisplayModeRotation.Unspecified:
                case DisplayModeRotation.Identity:
                    break;
                case DisplayModeRotation.Rotate90:
                    bitmap.RotateFlip(RotateFlipType.Rotate270FlipNone);
                    break;
                case DisplayModeRotation.Rotate180:
                    bitmap.RotateFlip(RotateFlipType.Rotate180FlipNone);
                    break;
                case DisplayModeRotation.Rotate270:
                    bitmap.RotateFlip(RotateFlipType.Rotate90FlipNone);
                    break;
                default:
                    break;
            }

            return DxCaptureResult.Ok(bitmap.ToSKBitmap(), result);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error while getting DirectX frame.");
        }
        finally
        {
            try
            {
                dxOutput.OutputDuplication.ReleaseFrame();
            }
            catch { }
        }

        return DxCaptureResult.Fail("Failed to get DirectX frame.");
    }

    private void ClearDirectXOutputs()
    {
        foreach (var screen in _directxScreens.Values)
        {
            try
            {
                screen.Dispose();
            }
            catch { }
        }

        _directxScreens.Clear();
    }

    private void InitBitBlt()
    {
        _bitBltScreens.Clear();

        for (var i = 0; i < Screen.AllScreens.Length; i++)
        {
            _bitBltScreens.Add(Screen.AllScreens[i].DeviceName, i);
        }
    }

    private void InitDirectX()
    {
        try
        {
            ClearDirectXOutputs();

            using var factory = new Factory1();
            foreach (var adapter in factory.Adapters1.Where(x => (x.Outputs?.Length ?? 0) > 0))
            {
                foreach (var output in adapter.Outputs)
                {
                    //_logger.LogWarning(output.Description.DeviceName);

                    try
                    {
                        var device = new SharpDX.Direct3D11.Device(adapter);
                        var output1 = output.QueryInterface<Output1>();

                        var bounds = output1.Description.DesktopBounds;
                        var width = bounds.Right - bounds.Left;
                        var height = bounds.Bottom - bounds.Top;

                        // Create Staging texture CPU-accessible
                        var textureDesc = new Texture2DDescription
                        {
                            CpuAccessFlags = CpuAccessFlags.Read,
                            BindFlags = BindFlags.None,
                            Format = Format.B8G8R8A8_UNorm,
                            Width = width,
                            Height = height,
                            OptionFlags = ResourceOptionFlags.None,
                            MipLevels = 1,
                            ArraySize = 1,
                            SampleDescription = { Count = 1, Quality = 0 },
                            Usage = ResourceUsage.Staging
                        };

                        var texture2D = new Texture2D(device, textureDesc);

                        _directxScreens.Add(
                            output1.Description.DeviceName,
                            new DirectXOutput(adapter,
                                device,
                                output1.DuplicateOutput(device),
                                texture2D,
                                output1.Description.Rotation));
                    }
                    catch (Exception ex)
                    {
                        _logger.LogError(ex, "Error while initializing DirectX.");
                    }
                }
            }
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error while initializing DirectX.");
        }
    }

    private bool IsEmpty(SKBitmap bitmap)
    {
        if (bitmap is null) return true;

        var height = bitmap.Height;
        var width = bitmap.Width;
        var bytesPerPixel = bitmap.BytesPerPixel;

        try
        {
            unsafe
            {
                byte* scan = (byte*)bitmap.GetPixels();

                for (var row = 0; row < height; row++)
                {
                    for (var column = 0; column < width; column++)
                    {
                        var index = row * width * bytesPerPixel + column * bytesPerPixel;

                        byte* data = scan + index;

                        for (var i = 0; i < bytesPerPixel; i++)
                        {
                            if (data[i] != 0)
                            {
                                return false;
                            }
                        }
                    }
                }

                return true;
            }
        }
        catch
        {
            return true;
        }
    }

    private void RefreshCurrentScreenBounds()
    {
        CurrentScreenBounds = Screen.AllScreens[_bitBltScreens[SelectedScreen]].Bounds;

        _initialized = false;
    }

    private void SystemEvents_DisplaySettingsChanged(object? sender, EventArgs e)
    {
        RefreshCurrentScreenBounds();
    }

    public void Dispose()
    {
        try
        {
            SystemEvents.DisplaySettingsChanged -= SystemEvents_DisplaySettingsChanged;
            ClearDirectXOutputs();
            GC.SuppressFinalize(this);
        }
        catch { }
    }

    private class DxCaptureResult
    {
        public SKBitmap? Bitmap { get; init; }
        public SharpDX.Result? DxResult { get; init; }
        public string FailureReason { get; init; } = string.Empty;

        [MemberNotNull(nameof(Bitmap))]
        public bool HadChanges { get; init; }

        public bool IsSuccess { get; init; }

        internal static DxCaptureResult Fail(string failureReason)
        {
            return new DxCaptureResult()
            {
                FailureReason = failureReason
            };
        }

        internal static DxCaptureResult Fail(string failureReason, SharpDX.Result dxResult)
        {
            return new DxCaptureResult()
            {
                FailureReason = failureReason,
                DxResult = dxResult
            };
        }

        internal static DxCaptureResult NoAccumulatedFrames(SharpDX.Result dxResult)
        {
            return new DxCaptureResult()
            {
                FailureReason = "No frames were accumulated.",
                DxResult = dxResult,
                IsSuccess = true
            };
        }

        internal static DxCaptureResult Ok(SKBitmap sKBitmap, SharpDX.Result result)
        {
            return new DxCaptureResult()
            {
                Bitmap = sKBitmap,
                DxResult = result,
                HadChanges = true,
                IsSuccess = true,
            };
        }

        internal static DxCaptureResult TryAcquireFailed(SharpDX.Result dxResult)
        {
            if (dxResult.Code == SharpDX.DXGI.ResultCode.WaitTimeout.Code)
            {
                return new DxCaptureResult()
                {
                    FailureReason = "Timed out while waiting for the next frame.",
                    DxResult = dxResult,
                    IsSuccess = true
                };
            }
            return new DxCaptureResult()
            {
                FailureReason = "TryAcquireFrame returned failure.",
                DxResult = dxResult
            };
        }
    }
}