using System;
using System.Drawing;
using System.Drawing.Imaging;
using System.IO;
using System.Runtime.InteropServices;
using SharpDX;
using SharpDX.Direct3D11;
using SharpDX.DXGI;
using TurboJpegWrapper;
using Device = SharpDX.Direct3D11.Device;
using MapFlags = SharpDX.Direct3D11.MapFlags;
using Resource = SharpDX.DXGI.Resource;

namespace Crysome.Client.Handlers
{
    internal sealed class DxgiCapture : IDisposable
    {
        private Device _device;
        private OutputDuplication _duplication;
        private Texture2D _staging;
        private int _width;
        private int _height;
        private bool _disposed;

        private byte[] _bgrBuffer;
        private TJCompressor _turboJpeg;
        private bool _turboAvailable;

        public int Width => _width;
        public int Height => _height;
        public bool IsInitialized => _duplication != null;

        public bool Init(int screenIndex)
        {
            Cleanup();
            try
            {
                using (var factory = new Factory1())
                {
                    Adapter1 adapter = factory.GetAdapter1(0);
                    _device = new Device(adapter, DeviceCreationFlags.BgraSupport);

                    int outputIdx = Math.Max(0, screenIndex);
                    Output output = adapter.GetOutput(outputIdx);
                    using (var output1 = output.QueryInterface<Output1>())
                    {
                        _width = output.Description.DesktopBounds.Right - output.Description.DesktopBounds.Left;
                        _height = output.Description.DesktopBounds.Bottom - output.Description.DesktopBounds.Top;

                        _staging = new Texture2D(_device, new Texture2DDescription
                        {
                            CpuAccessFlags = CpuAccessFlags.Read,
                            BindFlags = BindFlags.None,
                            Format = Format.B8G8R8A8_UNorm,
                            Width = _width,
                            Height = _height,
                            MipLevels = 1,
                            ArraySize = 1,
                            SampleDescription = new SampleDescription(1, 0),
                            Usage = ResourceUsage.Staging,
                            OptionFlags = ResourceOptionFlags.None
                        });

                        _duplication = output1.DuplicateOutput(_device);
                    }
                    output.Dispose();
                    adapter.Dispose();
                }

                _bgrBuffer = new byte[_width * _height * 3];

                try
                {
                    _turboJpeg = new TJCompressor();
                    _turboAvailable = true;
                }
                catch
                {
                    _turboAvailable = false;
                }

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

        public byte[] CaptureFrameToJpeg(int quality, int timeoutMs = 8)
        {
            if (_duplication == null)
                return null;

            Resource desktopResource = null;
            try
            {
                var result = _duplication.TryAcquireNextFrame(timeoutMs, out _, out desktopResource);
                if (result.Failure || desktopResource == null)
                    return null;

                using (var texture = desktopResource.QueryInterface<Texture2D>())
                {
                    _device.ImmediateContext.CopyResource(texture, _staging);
                }

                var mapSource = _device.ImmediateContext.MapSubresource(_staging, 0, MapMode.Read, MapFlags.None);
                try
                {
                    ConvertBgraToRgb(mapSource.DataPointer, mapSource.RowPitch);
                }
                finally
                {
                    _device.ImmediateContext.UnmapSubresource(_staging, 0);
                }

                if (_turboAvailable)
                    return TurboEncode(quality);
                else
                    return GdiEncode(quality);
            }
            catch (SharpDXException ex) when (ex.ResultCode.Code == SharpDX.DXGI.ResultCode.AccessLost.Result.Code)
            {
                Cleanup();
                return null;
            }
            catch
            {
                return null;
            }
            finally
            {
                try { desktopResource?.Dispose(); } catch { }
                try { _duplication?.ReleaseFrame(); } catch { }
            }
        }

        private unsafe void ConvertBgraToRgb(IntPtr srcPtr, int srcStride)
        {
            fixed (byte* dstBase = _bgrBuffer)
            {
                byte* src = (byte*)srcPtr;
                byte* dst = dstBase;
                int dstStride = _width * 3;

                for (int y = 0; y < _height; y++)
                {
                    byte* srcRow = src + y * srcStride;
                    byte* dstRow = dst + y * dstStride;
                    int x = 0;

                    for (; x < _width; x++)
                    {
                        dstRow[0] = srcRow[2]; // R (turbo wants RGB)
                        dstRow[1] = srcRow[1]; // G
                        dstRow[2] = srcRow[0]; // B
                        srcRow += 4;
                        dstRow += 3;
                    }
                }
            }
        }

        private byte[] TurboEncode(int quality)
        {
            try
            {
                return _turboJpeg.Compress(
                    _bgrBuffer,
                    _width * 3,
                    _width,
                    _height,
                    TJPixelFormats.TJPF_RGB,
                    TJSubsamplingOptions.TJSAMP_420,
                    quality,
                    TJFlags.FASTDCT);
            }
            catch
            {
                _turboAvailable = false;
                return GdiEncode(quality);
            }
        }

        private byte[] GdiEncode(int quality)
        {
            using (var bmp = new Bitmap(_width, _height, PixelFormat.Format24bppRgb))
            {
                var bmpData = bmp.LockBits(
                    new Rectangle(0, 0, _width, _height),
                    ImageLockMode.WriteOnly,
                    PixelFormat.Format24bppRgb);

                int bgrStride = _width * 3;
                int dstStride = bmpData.Stride;

                unsafe
                {
                    fixed (byte* srcBase = _bgrBuffer)
                    {
                        byte* dst = (byte*)bmpData.Scan0;
                        for (int y = 0; y < _height; y++)
                        {
                            // BGR buffer has RGB order, but Bitmap Format24bppRgb expects BGR
                            byte* srcRow = srcBase + y * bgrStride;
                            byte* dstRow = dst + y * dstStride;
                            for (int x = 0; x < _width; x++)
                            {
                                dstRow[0] = srcRow[2]; // B
                                dstRow[1] = srcRow[1]; // G
                                dstRow[2] = srcRow[0]; // R
                                srcRow += 3;
                                dstRow += 3;
                            }
                        }
                    }
                }

                bmp.UnlockBits(bmpData);

                var codec = GetJpegCodec();
                using (var ms = new MemoryStream(131072))
                {
                    if (codec != null)
                    {
                        using (var ep = new EncoderParameters(1))
                        {
                            ep.Param[0] = new EncoderParameter(Encoder.Quality, quality);
                            bmp.Save(ms, codec, ep);
                        }
                    }
                    else
                    {
                        bmp.Save(ms, ImageFormat.Jpeg);
                    }
                    return ms.ToArray();
                }
            }
        }

        private static ImageCodecInfo _jpegCodec;
        private static readonly object _codecLock = new object();

        private static ImageCodecInfo GetJpegCodec()
        {
            if (_jpegCodec != null) return _jpegCodec;
            lock (_codecLock)
            {
                if (_jpegCodec != null) return _jpegCodec;
                foreach (var info in ImageCodecInfo.GetImageEncoders())
                    if (info.FormatID == ImageFormat.Jpeg.Guid)
                    { _jpegCodec = info; break; }
            }
            return _jpegCodec;
        }

        private void Cleanup()
        {
            try { _turboJpeg?.Dispose(); } catch { }
            try { _duplication?.Dispose(); } catch { }
            try { _staging?.Dispose(); } catch { }
            try { _device?.Dispose(); } catch { }
            _turboJpeg = null;
            _duplication = null;
            _staging = null;
            _device = null;
            _bgrBuffer = null;
        }

        public void Dispose()
        {
            if (!_disposed)
            {
                _disposed = true;
                Cleanup();
            }
        }
    }
}