/****************************************************************************
 * Copyright (C) 2016 Intel Corporation.   All Rights Reserved.
 *
 * 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 (including the next
 * paragraph) 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.
 *
 * @file archrast.cpp
 *
 * @brief Implementation for archrast.
 *
 ******************************************************************************/
#include <atomic>

#include "common/os.h"
#include "archrast/archrast.h"
#include "archrast/eventmanager.h"
#include "gen_ar_eventhandlerfile.hpp"

namespace ArchRast
{
    //////////////////////////////////////////////////////////////////////////
    /// @brief struct that keeps track of depth and stencil event information
    struct DepthStencilStats
    {
        uint32_t earlyZTestPassCount       = 0;
        uint32_t earlyZTestFailCount       = 0;
        uint32_t lateZTestPassCount        = 0;
        uint32_t lateZTestFailCount        = 0;
        uint32_t earlyStencilTestPassCount = 0;
        uint32_t earlyStencilTestFailCount = 0;
        uint32_t lateStencilTestPassCount  = 0;
        uint32_t lateStencilTestFailCount  = 0;
    };

    struct CStats
    {
        uint32_t trivialRejectCount;
        uint32_t trivialAcceptCount;
        uint32_t mustClipCount;
    };

    struct TEStats
    {
        uint32_t inputPrims = 0;
        //@todo:: Change this to numPatches. Assumed: 1 patch per prim. If holds, its fine.
    };

    struct GSStateInfo
    {
        uint32_t inputPrimCount;
        uint32_t primGeneratedCount;
        uint32_t vertsInput;
    };

    struct RastStats
    {
        uint32_t rasterTiles = 0;
    };

    struct CullStats
    {
        uint32_t degeneratePrimCount = 0;
        uint32_t backfacePrimCount   = 0;
    };

    struct AlphaStats
    {
        uint32_t alphaTestCount  = 0;
        uint32_t alphaBlendCount = 0;
    };

    //////////////////////////////////////////////////////////////////////////
    /// @brief Event handler that handles API thread events. This is shared
    ///        between the API and its caller (e.g. driver shim) but typically
    ///        there is only a single API thread per context. So you can save
    ///        information in the class to be used for other events.
    class EventHandlerApiStats : public EventHandlerFile
    {
    public:
        EventHandlerApiStats(uint32_t id) : EventHandlerFile(id)
        {
#if defined(_WIN32)
            // Attempt to copy the events.proto file to the ArchRast output dir. It's common for
            // tools to place the events.proto file in the DEBUG_OUTPUT_DIR when launching AR. If it
            // exists, this will attempt to copy it the first time we get here to package it with
            // the stats. Otherwise, the user would need to specify the events.proto location when
            // parsing the stats in post.
            std::stringstream eventsProtoSrcFilename, eventsProtoDstFilename;
            eventsProtoSrcFilename << KNOB_DEBUG_OUTPUT_DIR << "\\events.proto" << std::ends;
            eventsProtoDstFilename << mOutputDir.substr(0, mOutputDir.size() - 1)
                                   << "\\events.proto" << std::ends;

            // If event.proto already exists, we're done; else do the copy
            struct stat buf; // Use a Posix stat for file existence check
            if (!stat(eventsProtoDstFilename.str().c_str(), &buf) == 0)
            {
                // Now check to make sure the events.proto source exists
                if (stat(eventsProtoSrcFilename.str().c_str(), &buf) == 0)
                {
                    std::ifstream srcFile;
                    srcFile.open(eventsProtoSrcFilename.str().c_str(), std::ios::binary);
                    if (srcFile.is_open())
                    {
                        // Just do a binary buffer copy
                        std::ofstream dstFile;
                        dstFile.open(eventsProtoDstFilename.str().c_str(), std::ios::binary);
                        dstFile << srcFile.rdbuf();
                        dstFile.close();
                    }
                    srcFile.close();
                }
            }
#endif
        }

        virtual void Handle(const DrawInstancedEvent& event)
        {
            DrawInfoEvent e(event.data.drawId,
                            ArchRast::Instanced,
                            event.data.topology,
                            event.data.numVertices,
                            0,
                            0,
                            event.data.startVertex,
                            event.data.numInstances,
                            event.data.startInstance,
                            event.data.tsEnable,
                            event.data.gsEnable,
                            event.data.soEnable,
                            event.data.soTopology,
                            event.data.splitId);

            EventHandlerFile::Handle(e);
        }

        virtual void Handle(const DrawIndexedInstancedEvent& event)
        {
            DrawInfoEvent e(event.data.drawId,
                            ArchRast::IndexedInstanced,
                            event.data.topology,
                            0,
                            event.data.numIndices,
                            event.data.indexOffset,
                            event.data.baseVertex,
                            event.data.numInstances,
                            event.data.startInstance,
                            event.data.tsEnable,
                            event.data.gsEnable,
                            event.data.soEnable,
                            event.data.soTopology,
                            event.data.splitId);

            EventHandlerFile::Handle(e);
        }
    };

    //////////////////////////////////////////////////////////////////////////
    /// @brief Event handler that handles worker thread events. There is one
    ///        event handler per thread. The python script will need to sum
    ///        up counters across all of the threads.
    class EventHandlerWorkerStats : public EventHandlerFile
    {
    public:
        EventHandlerWorkerStats(uint32_t id) : EventHandlerFile(id), mNeedFlush(false)
        {
            memset(mShaderStats, 0, sizeof(mShaderStats));
        }

        virtual void Handle(const EarlyDepthStencilInfoSingleSample& event)
        {
            // earlyZ test compute
            mDSSingleSample.earlyZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSSingleSample.earlyZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // earlyStencil test compute
            mDSSingleSample.earlyStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSSingleSample.earlyStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);

            // earlyZ test single and multi sample
            mDSCombined.earlyZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSCombined.earlyZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // earlyStencil test single and multi sample
            mDSCombined.earlyStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSCombined.earlyStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);

            mNeedFlush = true;
        }

        virtual void Handle(const EarlyDepthStencilInfoSampleRate& event)
        {
            // earlyZ test compute
            mDSSampleRate.earlyZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSSampleRate.earlyZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // earlyStencil test compute
            mDSSampleRate.earlyStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSSampleRate.earlyStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);

            // earlyZ test single and multi sample
            mDSCombined.earlyZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSCombined.earlyZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // earlyStencil test single and multi sample
            mDSCombined.earlyStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSCombined.earlyStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);

            mNeedFlush = true;
        }

        virtual void Handle(const EarlyDepthStencilInfoNullPS& event)
        {
            // earlyZ test compute
            mDSNullPS.earlyZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSNullPS.earlyZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // earlyStencil test compute
            mDSNullPS.earlyStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSNullPS.earlyStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);
            mNeedFlush = true;
        }

        virtual void Handle(const LateDepthStencilInfoSingleSample& event)
        {
            // lateZ test compute
            mDSSingleSample.lateZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSSingleSample.lateZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // lateStencil test compute
            mDSSingleSample.lateStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSSingleSample.lateStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);

            // lateZ test single and multi sample
            mDSCombined.lateZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSCombined.lateZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // lateStencil test single and multi sample
            mDSCombined.lateStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSCombined.lateStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);

            mNeedFlush = true;
        }

        virtual void Handle(const LateDepthStencilInfoSampleRate& event)
        {
            // lateZ test compute
            mDSSampleRate.lateZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSSampleRate.lateZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // lateStencil test compute
            mDSSampleRate.lateStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSSampleRate.lateStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);

            // lateZ test single and multi sample
            mDSCombined.lateZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSCombined.lateZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // lateStencil test single and multi sample
            mDSCombined.lateStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSCombined.lateStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);

            mNeedFlush = true;
        }

        virtual void Handle(const LateDepthStencilInfoNullPS& event)
        {
            // lateZ test compute
            mDSNullPS.lateZTestPassCount += _mm_popcnt_u32(event.data.depthPassMask);
            mDSNullPS.lateZTestFailCount +=
                _mm_popcnt_u32((!event.data.depthPassMask) & event.data.coverageMask);

            // lateStencil test compute
            mDSNullPS.lateStencilTestPassCount += _mm_popcnt_u32(event.data.stencilPassMask);
            mDSNullPS.lateStencilTestFailCount +=
                _mm_popcnt_u32((!event.data.stencilPassMask) & event.data.coverageMask);
            mNeedFlush = true;
        }

        virtual void Handle(const EarlyDepthInfoPixelRate& event)
        {
            // earlyZ test compute
            mDSPixelRate.earlyZTestPassCount += event.data.depthPassCount;
            mDSPixelRate.earlyZTestFailCount +=
                (_mm_popcnt_u32(event.data.activeLanes) - event.data.depthPassCount);
            mNeedFlush = true;
        }


        virtual void Handle(const LateDepthInfoPixelRate& event)
        {
            // lateZ test compute
            mDSPixelRate.lateZTestPassCount += event.data.depthPassCount;
            mDSPixelRate.lateZTestFailCount +=
                (_mm_popcnt_u32(event.data.activeLanes) - event.data.depthPassCount);
            mNeedFlush = true;
        }


        virtual void Handle(const ClipInfoEvent& event)
        {
            mClipper.mustClipCount += _mm_popcnt_u32(event.data.clipMask);
            mClipper.trivialRejectCount +=
                event.data.numInvocations - _mm_popcnt_u32(event.data.validMask);
            mClipper.trivialAcceptCount +=
                _mm_popcnt_u32(event.data.validMask & ~event.data.clipMask);
        }

        struct ShaderStats
        {
            uint32_t numInstExecuted;
        };

        virtual void Handle(const VSStats& event)
        {
            mShaderStats[SHADER_VERTEX].numInstExecuted += event.data.numInstExecuted;
        }

        virtual void Handle(const GSStats& event)
        {
            mShaderStats[SHADER_GEOMETRY].numInstExecuted += event.data.numInstExecuted;
        }

        virtual void Handle(const DSStats& event)
        {
            mShaderStats[SHADER_DOMAIN].numInstExecuted += event.data.numInstExecuted;
        }

        virtual void Handle(const HSStats& event)
        {
            mShaderStats[SHADER_HULL].numInstExecuted += event.data.numInstExecuted;
        }

        virtual void Handle(const PSStats& event)
        {
            mShaderStats[SHADER_PIXEL].numInstExecuted += event.data.numInstExecuted;
            mNeedFlush = true;
        }

        virtual void Handle(const CSStats& event)
        {
            mShaderStats[SHADER_COMPUTE].numInstExecuted += event.data.numInstExecuted;
            mNeedFlush = true;
        }

        // Flush cached events for this draw
        virtual void FlushDraw(uint32_t drawId)
        {
            if (mNeedFlush == false)
                return;

            EventHandlerFile::Handle(PSInfo(drawId, mShaderStats[SHADER_PIXEL].numInstExecuted));
            EventHandlerFile::Handle(CSInfo(drawId, mShaderStats[SHADER_COMPUTE].numInstExecuted));

            // singleSample
            EventHandlerFile::Handle(EarlyZSingleSample(
                drawId, mDSSingleSample.earlyZTestPassCount, mDSSingleSample.earlyZTestFailCount));
            EventHandlerFile::Handle(LateZSingleSample(
                drawId, mDSSingleSample.lateZTestPassCount, mDSSingleSample.lateZTestFailCount));
            EventHandlerFile::Handle(
                EarlyStencilSingleSample(drawId,
                                         mDSSingleSample.earlyStencilTestPassCount,
                                         mDSSingleSample.earlyStencilTestFailCount));
            EventHandlerFile::Handle(
                LateStencilSingleSample(drawId,
                                        mDSSingleSample.lateStencilTestPassCount,
                                        mDSSingleSample.lateStencilTestFailCount));

            // sampleRate
            EventHandlerFile::Handle(EarlyZSampleRate(
                drawId, mDSSampleRate.earlyZTestPassCount, mDSSampleRate.earlyZTestFailCount));
            EventHandlerFile::Handle(LateZSampleRate(
                drawId, mDSSampleRate.lateZTestPassCount, mDSSampleRate.lateZTestFailCount));
            EventHandlerFile::Handle(
                EarlyStencilSampleRate(drawId,
                                       mDSSampleRate.earlyStencilTestPassCount,
                                       mDSSampleRate.earlyStencilTestFailCount));
            EventHandlerFile::Handle(LateStencilSampleRate(drawId,
                                                           mDSSampleRate.lateStencilTestPassCount,
                                                           mDSSampleRate.lateStencilTestFailCount));

            // combined
            EventHandlerFile::Handle(
                EarlyZ(drawId, mDSCombined.earlyZTestPassCount, mDSCombined.earlyZTestFailCount));
            EventHandlerFile::Handle(
                LateZ(drawId, mDSCombined.lateZTestPassCount, mDSCombined.lateZTestFailCount));
            EventHandlerFile::Handle(EarlyStencil(drawId,
                                                  mDSCombined.earlyStencilTestPassCount,
                                                  mDSCombined.earlyStencilTestFailCount));
            EventHandlerFile::Handle(LateStencil(drawId,
                                                 mDSCombined.lateStencilTestPassCount,
                                                 mDSCombined.lateStencilTestFailCount));

            // pixelRate
            EventHandlerFile::Handle(EarlyZPixelRate(
                drawId, mDSPixelRate.earlyZTestPassCount, mDSPixelRate.earlyZTestFailCount));
            EventHandlerFile::Handle(LateZPixelRate(
                drawId, mDSPixelRate.lateZTestPassCount, mDSPixelRate.lateZTestFailCount));


            // NullPS
            EventHandlerFile::Handle(
                EarlyZNullPS(drawId, mDSNullPS.earlyZTestPassCount, mDSNullPS.earlyZTestFailCount));
            EventHandlerFile::Handle(EarlyStencilNullPS(
                drawId, mDSNullPS.earlyStencilTestPassCount, mDSNullPS.earlyStencilTestFailCount));

            // Rasterized Subspans
            EventHandlerFile::Handle(RasterTiles(drawId, rastStats.rasterTiles));

            // Alpha Subspans
            EventHandlerFile::Handle(
                AlphaEvent(drawId, mAlphaStats.alphaTestCount, mAlphaStats.alphaBlendCount));

            // Primitive Culling
            EventHandlerFile::Handle(
                CullEvent(drawId, mCullStats.backfacePrimCount, mCullStats.degeneratePrimCount));

            mDSSingleSample = {};
            mDSSampleRate   = {};
            mDSCombined     = {};
            mDSPixelRate    = {};
            mDSNullPS = {};

            rastStats   = {};
            mCullStats  = {};
            mAlphaStats = {};

            mShaderStats[SHADER_PIXEL]   = {};
            mShaderStats[SHADER_COMPUTE] = {};

            mNeedFlush = false;
        }

        virtual void Handle(const FrontendDrawEndEvent& event)
        {
            // Clipper
            EventHandlerFile::Handle(ClipperEvent(event.data.drawId,
                                                  mClipper.trivialRejectCount,
                                                  mClipper.trivialAcceptCount,
                                                  mClipper.mustClipCount));

            // Tesselator
            EventHandlerFile::Handle(TessPrims(event.data.drawId, mTS.inputPrims));

            // Geometry Shader
            EventHandlerFile::Handle(GSInputPrims(event.data.drawId, mGS.inputPrimCount));
            EventHandlerFile::Handle(GSPrimsGen(event.data.drawId, mGS.primGeneratedCount));
            EventHandlerFile::Handle(GSVertsInput(event.data.drawId, mGS.vertsInput));

            EventHandlerFile::Handle(
                VSInfo(event.data.drawId, mShaderStats[SHADER_VERTEX].numInstExecuted));
            EventHandlerFile::Handle(
                HSInfo(event.data.drawId, mShaderStats[SHADER_HULL].numInstExecuted));
            EventHandlerFile::Handle(
                DSInfo(event.data.drawId, mShaderStats[SHADER_DOMAIN].numInstExecuted));
            EventHandlerFile::Handle(
                GSInfo(event.data.drawId, mShaderStats[SHADER_GEOMETRY].numInstExecuted));

            mShaderStats[SHADER_VERTEX]   = {};
            mShaderStats[SHADER_HULL]     = {};
            mShaderStats[SHADER_DOMAIN]   = {};
            mShaderStats[SHADER_GEOMETRY] = {};

            // Reset Internal Counters
            mClipper = {};
            mTS      = {};
            mGS      = {};
        }

        virtual void Handle(const GSPrimInfo& event)
        {
            mGS.inputPrimCount += event.data.inputPrimCount;
            mGS.primGeneratedCount += event.data.primGeneratedCount;
            mGS.vertsInput += event.data.vertsInput;
        }

        virtual void Handle(const TessPrimCount& event) { mTS.inputPrims += event.data.primCount; }

        virtual void Handle(const RasterTileCount& event)
        {
            rastStats.rasterTiles += event.data.rasterTiles;
        }

        virtual void Handle(const CullInfoEvent& event)
        {
            mCullStats.degeneratePrimCount += _mm_popcnt_u32(
                event.data.validMask ^ (event.data.validMask & ~event.data.degeneratePrimMask));
            mCullStats.backfacePrimCount += _mm_popcnt_u32(
                event.data.validMask ^ (event.data.validMask & ~event.data.backfacePrimMask));
        }

        virtual void Handle(const AlphaInfoEvent& event)
        {
            mAlphaStats.alphaTestCount += event.data.alphaTestEnable;
            mAlphaStats.alphaBlendCount += event.data.alphaBlendEnable;
        }

    protected:
        bool mNeedFlush;
        // Per draw stats
        DepthStencilStats mDSSingleSample = {};
        DepthStencilStats mDSSampleRate   = {};
        DepthStencilStats mDSPixelRate    = {};
        DepthStencilStats mDSCombined     = {};
        DepthStencilStats mDSNullPS       = {};
        DepthStencilStats mDSOmZ          = {};
        CStats            mClipper        = {};
        TEStats           mTS             = {};
        GSStateInfo       mGS             = {};
        RastStats         rastStats       = {};
        CullStats         mCullStats      = {};
        AlphaStats        mAlphaStats     = {};

        ShaderStats mShaderStats[NUM_SHADER_TYPES];

    };

    static EventManager* FromHandle(HANDLE hThreadContext)
    {
        return reinterpret_cast<EventManager*>(hThreadContext);
    }

    // Construct an event manager and associate a handler with it.
    HANDLE CreateThreadContext(AR_THREAD type)
    {
        // Can we assume single threaded here?
        static std::atomic<uint32_t> counter(0);
        uint32_t                     id = counter.fetch_add(1);

        EventManager* pManager = new EventManager();

        if (pManager)
        {
            EventHandlerFile* pHandler = nullptr;

            if (type == AR_THREAD::API)
            {
                pHandler = new EventHandlerApiStats(id);
                pManager->Attach(pHandler);
                pHandler->Handle(ThreadStartApiEvent());
            }
            else
            {
                pHandler = new EventHandlerWorkerStats(id);
                pManager->Attach(pHandler);
                pHandler->Handle(ThreadStartWorkerEvent());
            }

            pHandler->MarkHeader();

            return pManager;
        }

        SWR_INVALID("Failed to register thread.");
        return nullptr;
    }

    void DestroyThreadContext(HANDLE hThreadContext)
    {
        EventManager* pManager = FromHandle(hThreadContext);
        SWR_ASSERT(pManager != nullptr);

        delete pManager;
    }

    // Dispatch event for this thread.
    void Dispatch(HANDLE hThreadContext, const Event& event)
    {
        EventManager* pManager = FromHandle(hThreadContext);
        SWR_ASSERT(pManager != nullptr);

        pManager->Dispatch(event);
    }

    // Flush for this thread.
    void FlushDraw(HANDLE hThreadContext, uint32_t drawId)
    {
        EventManager* pManager = FromHandle(hThreadContext);
        SWR_ASSERT(pManager != nullptr);

        pManager->FlushDraw(drawId);
    }
} // namespace ArchRast
