using System.Runtime.InteropServices;
using Nerfed.Runtime.Video;
using RefreshCS;

namespace Nerfed.Runtime.Graphics;

/// <summary>
/// Manages all graphics-related concerns.
/// </summary>
public class GraphicsDevice : IDisposable
{
    public IntPtr Handle { get; }
    public BackendFlags Backend { get; }
    public bool DebugMode { get; }

    // Built-in video pipeline
    internal GraphicsPipeline VideoPipeline { get; }

    // Built-in text shader info
    public Shader TextVertexShader;
    public Shader TextFragmentShader;
    public VertexInputState TextVertexInputState { get; }

    // Built-in samplers
    public Sampler PointSampler { get; }
    public Sampler LinearSampler { get; }

    public bool IsDisposed { get; private set; }

    internal readonly RenderPassPool RenderPassPool = new RenderPassPool();
    internal readonly ComputePassPool ComputePassPool = new ComputePassPool();
    internal readonly CopyPassPool CopyPassPool = new CopyPassPool();
    private readonly HashSet<GCHandle> resources = new HashSet<GCHandle>();
    private readonly CommandBufferPool commandBufferPool;
    private readonly FencePool fencePool;

    internal GraphicsDevice(BackendFlags preferredBackends)
    {
        if (preferredBackends == BackendFlags.Invalid)
        {
            throw new System.Exception("Could not set graphics backend!");
        }

        bool debugMode = false;
#if DEBUG
        debugMode = true;
#endif

        Handle = Refresh.Refresh_CreateDevice((Refresh.BackendFlags)preferredBackends, Conversions.BoolToInt(debugMode));

        DebugMode = debugMode;
        // TODO: check for CreateDevice fail

        Backend = (BackendFlags)Refresh.Refresh_GetBackend(Handle);

        IEmbeddedShaders embeddedShaders;
        switch (Backend)
        {
            case BackendFlags.Vulkan:
                embeddedShaders = new EmbeddedShadersSpirV();
                break;
            case BackendFlags.D3D11:
                throw new NotImplementedException("D3D11 embedded shaders");
                break;
            case BackendFlags.Metal:
                throw new NotImplementedException("Metal embedded shaders");
                break;
            default: throw new ArgumentOutOfRangeException();
        }

        Shader fullscreenVertShader;
        Shader textVertShader;
        Shader textFragShader;
        Shader videoFragShader;

        using (MemoryStream fullscreenVertStream = new MemoryStream(embeddedShaders.FullscreenVert))
        {
            fullscreenVertShader = new Shader(
                this,
                fullscreenVertStream,
                "main",
                new ShaderCreateInfo
                {
                    ShaderStage = ShaderStage.Vertex,
                    ShaderFormat = embeddedShaders.ShaderFormat
                }
            );
        }

        using (MemoryStream videoYuv2RgbaFragStream = new MemoryStream(embeddedShaders.VideoYuv2RgbaFrag))
        {
            videoFragShader = new Shader(
                this,
                videoYuv2RgbaFragStream,
                "main",
                new ShaderCreateInfo
                {
                    ShaderStage = ShaderStage.Fragment,
                    ShaderFormat = embeddedShaders.ShaderFormat,
                    SamplerCount = 3
                }
            );
        }

        using (MemoryStream textTransformVertStream = new MemoryStream(embeddedShaders.TextTransformVert))
        {
            textVertShader = new Shader(
                this,
                textTransformVertStream,
                "main",
                new ShaderCreateInfo
                {
                    ShaderStage = ShaderStage.Vertex,
                    ShaderFormat = embeddedShaders.ShaderFormat,
                    UniformBufferCount = 1
                }
            );
        }

        using (MemoryStream textMsdfFragStream = new MemoryStream(embeddedShaders.TextMsdfFrag))
        {
            textFragShader = new Shader(
                this,
                textMsdfFragStream,
                "main",
                new ShaderCreateInfo
                {
                    ShaderStage = ShaderStage.Fragment,
                    ShaderFormat = embeddedShaders.ShaderFormat,
                    SamplerCount = 1,
                    UniformBufferCount = 1
                }
            );
        }

        VideoPipeline = new GraphicsPipeline(
            this,
            new GraphicsPipelineCreateInfo
            {
                AttachmentInfo = new GraphicsPipelineAttachmentInfo(
                    new ColorAttachmentDescription(
                        TextureFormat.R8G8B8A8,
                        ColorAttachmentBlendState.None
                    )
                ),
                DepthStencilState = DepthStencilState.Disable,
                VertexShader = fullscreenVertShader,
                FragmentShader = videoFragShader,
                VertexInputState = VertexInputState.Empty,
                RasterizerState = RasterizerState.CCW_CullNone,
                PrimitiveType = PrimitiveType.TriangleList,
                MultisampleState = MultisampleState.None
            }
        );

        TextVertexShader = textVertShader;
        TextFragmentShader = textFragShader;

        TextVertexInputState = VertexInputState.CreateSingleBinding<FontVertex>();

        PointSampler = new Sampler(this, SamplerCreateInfo.PointClamp);
        LinearSampler = new Sampler(this, SamplerCreateInfo.LinearClamp);

        fencePool = new FencePool(this);
        commandBufferPool = new CommandBufferPool(this);
    }

    /// <summary>
    /// Prepares a window so that frames can be presented to it.
    /// </summary>
    /// <param name="swapchainComposition">The desired composition of the swapchain. Ignore this unless you are using HDR or tonemapping.</param>
    /// <param name="presentMode">The desired presentation mode for the window. Roughly equivalent to V-Sync.</param>
    /// <returns>True if successfully claimed.</returns>
    public bool ClaimWindow(
        Window window,
        SwapchainComposition swapchainComposition,
        PresentMode presentMode
    )
    {
        if (window.Claimed)
        {
            Log.Error("Window already claimed!");
            return false;
        }

        bool success = Conversions.IntToBool(
            Refresh.Refresh_ClaimWindow(
                Handle,
                window.Handle,
                (Refresh.SwapchainComposition)swapchainComposition,
                (Refresh.PresentMode)presentMode
            )
        );

        if (success)
        {
            window.Claimed = true;
            window.SwapchainComposition = swapchainComposition;
            window.SwapchainFormat = GetSwapchainFormat(window);

            if (window.SwapchainTexture == null)
            {
                window.SwapchainTexture = new Texture(this, window.SwapchainFormat);
            }
        }

        return success;
    }

    /// <summary>
    /// Unclaims a window, making it unavailable for presenting and freeing associated resources.
    /// </summary>
    public void UnclaimWindow(Window window)
    {
        if (window.Claimed)
        {
            Refresh.Refresh_UnclaimWindow(
                Handle,
                window.Handle
            );
            window.Claimed = false;

            // The swapchain texture doesn't actually have a permanent texture reference, so we zero the handle before disposing.
            window.SwapchainTexture.Handle = IntPtr.Zero;
            window.SwapchainTexture.Dispose();
            window.SwapchainTexture = null;
        }
    }

    /// <summary>
    /// Changes the present mode of a claimed window. Does nothing if the window is not claimed.
    /// </summary>
    public bool SetSwapchainParameters(
        Window window,
        SwapchainComposition swapchainComposition,
        PresentMode presentMode
    )
    {
        if (!window.Claimed)
        {
            Log.Error("Cannot set present mode on unclaimed window!");
            return false;
        }

        bool success = Conversions.IntToBool(
            Refresh.Refresh_SetSwapchainParameters(
                Handle,
                window.Handle,
                (Refresh.SwapchainComposition)swapchainComposition,
                (Refresh.PresentMode)presentMode
            )
        );

        if (success)
        {
            window.SwapchainComposition = swapchainComposition;
            window.SwapchainFormat = GetSwapchainFormat(window);

            if (window.SwapchainTexture != null)
            {
                window.SwapchainTexture.Format = window.SwapchainFormat;
            }
        }

        return success;
    }

    /// <summary>
    /// Acquires a command buffer.
    /// This is the start of your rendering process.
    /// </summary>
    /// <returns></returns>
    public CommandBuffer AcquireCommandBuffer()
    {
        CommandBuffer commandBuffer = commandBufferPool.Obtain();
        commandBuffer.SetHandle(Refresh.Refresh_AcquireCommandBuffer(Handle));
#if DEBUG
        commandBuffer.ResetStateTracking();
#endif
        return commandBuffer;
    }

    /// <summary>
    /// Submits a command buffer to the GPU for processing.
    /// </summary>
    public void Submit(CommandBuffer commandBuffer)
    {
#if DEBUG
        if (commandBuffer.Submitted)
        {
            throw new System.InvalidOperationException("Command buffer already submitted!");
        }
#endif

        Refresh.Refresh_Submit(
            commandBuffer.Handle
        );

        commandBufferPool.Return(commandBuffer);

#if DEBUG
        commandBuffer.Submitted = true;
#endif
    }

    /// <summary>
    /// Submits a command buffer to the GPU for processing and acquires a fence associated with the submission.
    /// </summary>
    /// <returns></returns>
    public Fence SubmitAndAcquireFence(CommandBuffer commandBuffer)
    {
        IntPtr fenceHandle = Refresh.Refresh_SubmitAndAcquireFence(
            commandBuffer.Handle
        );

        Fence fence = fencePool.Obtain();
        fence.SetHandle(fenceHandle);

        return fence;
    }

    /// <summary>
    /// Wait for the graphics device to become idle.
    /// </summary>
    public void Wait()
    {
        Refresh.Refresh_Wait(Handle);
    }

    /// <summary>
    /// Waits for the given fence to become signaled.
    /// </summary>
    public unsafe void WaitForFence(Fence fence)
    {
        IntPtr fenceHandle = fence.Handle;

        Refresh.Refresh_WaitForFences(
            Handle,
            1,
            &fenceHandle,
            1
        );
    }

    /// <summary>
    /// Wait for one or more fences to become signaled.
    /// </summary>
    /// <param name="waitAll">If true, will wait for all given fences to be signaled.</param>
    public unsafe void WaitForFences(Span<Fence> fences, bool waitAll)
    {
        IntPtr* handlePtr = stackalloc nint[fences.Length];

        for (int i = 0; i < fences.Length; i += 1)
        {
            handlePtr[i] = fences[i].Handle;
        }

        Refresh.Refresh_WaitForFences(
            Handle,
            Conversions.BoolToInt(waitAll),
            handlePtr,
            (uint)fences.Length
        );
    }

    /// <summary>
    /// Returns true if the fence is signaled, indicating that the associated command buffer has finished processing.
    /// </summary>
    /// <exception cref="InvalidOperationException">Throws if the fence query indicates that the graphics device has been lost.</exception>
    public bool QueryFence(Fence fence)
    {
        int result = Refresh.Refresh_QueryFence(Handle, fence.Handle);

        if (result < 0)
        {
            throw new InvalidOperationException("The graphics device has been lost.");
        }

        return result != 0;
    }

    /// <summary>
    /// Release reference to an acquired fence, enabling it to be reused.
    /// </summary>
    public void ReleaseFence(Fence fence)
    {
        Refresh.Refresh_ReleaseFence(Handle, fence.Handle);
        fence.Handle = IntPtr.Zero;
        fencePool.Return(fence);
    }

    private TextureFormat GetSwapchainFormat(Window window)
    {
        if (!window.Claimed)
        {
            throw new System.ArgumentException("Cannot get swapchain format of unclaimed window!");
        }

        return (TextureFormat)Refresh.Refresh_GetSwapchainTextureFormat(Handle, window.Handle);
    }

    internal void AddResourceReference(GCHandle resourceReference)
    {
        lock (resources)
        {
            resources.Add(resourceReference);
        }
    }

    internal void RemoveResourceReference(GCHandle resourceReference)
    {
        lock (resources)
        {
            resources.Remove(resourceReference);
        }
    }

    private void Dispose(bool disposing)
    {
        if (!IsDisposed)
        {
            if (disposing)
            {
                lock (resources)
                {
                    // Dispose video players first to avoid race condition on threaded decoding
                    foreach (GCHandle resource in resources)
                    {
                        if (resource.Target is VideoPlayer player)
                        {
                            player.Dispose();
                        }
                    }

                    // Dispose everything else
                    foreach (GCHandle resource in resources)
                    {
                        if (resource.Target is IDisposable disposable)
                        {
                            disposable.Dispose();
                        }
                    }

                    resources.Clear();
                }
            }

            Refresh.Refresh_DestroyDevice(Handle);

            IsDisposed = true;
        }
    }

    ~GraphicsDevice()
    {
        // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
        Dispose(disposing: false);
    }

    public void Dispose()
    {
        // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
        Dispose(disposing: true);
        GC.SuppressFinalize(this);
    }
}