/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: OSThread.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "Threads.hpp"
#include "OSThread.hpp"
#include "../WaitFor.hpp"
#include "ThreadHandles.hpp"
#include "TLSView.hpp"

#if defined(AURORA_IS_LINUX_DERIVED)
    #include <sys/resource.h>
#endif

#if defined(AURORA_HAS_PTHREADS)
    #include <sched.h>
#endif

#if defined(AURORA_PLATFORM_WIN32)
    #include <Aux_ulib.h>
#endif

#include "SpawnThread.hpp"

namespace Aurora::Threading::Threads
{
    static ThreadInfo gDummyThreadInfo;

    OSThread::OSThread(const ThreadInfo &info) : info_(info)
    {
        name_ = info.name.value_or("Aurora Thread");
        terminated_ = Primitives::EventShared(true, false);

        // maybe we should atomic exchange compare these when needed frogthink
        terminateSignal_    = Primitives::EventShared(true, false, true);
        terminatedSignalLs_ = Loop::NewLSEvent(true, false);
        terminateSignalLs_  = Loop::NewLSEvent(true, false, true);

        exitOnlyOnce_ = Primitives::CriticalSectionUnique();
        SysAssert(terminated_ ? true : false, "out of memory");
    }

    OSThread::OSThread() : info_(gDummyThreadInfo)
    {
        name_ = "Main Thread";
        terminated_ = Primitives::EventShared(true, false);
        exitOnlyOnce_ = Primitives::CriticalSectionUnique();
    }

    OSThread::OSThread(AuUInt64 id) : info_(gDummyThreadInfo)
    {
        name_ = "System Thread";
        handle_ = reinterpret_cast<decltype(handle_)>(id);
    }

    OSThread::~OSThread()
    {
        bool bDetached {};
        bool bDetachedSuccess {};

        if (contextUsed_)
        {
            if (detached_)
            {
                bDetached = true;
            }
            else
            {
                Exit();
                WaitFor(terminated_.get());
            }
        }

        terminated_.reset();
        FreeOSContext();

        if (bDetached)
        {
            if (this->tlsReferenceThread_)
            {
                AU_LOCK_GUARD(this->tlsLock_);
                AU_LOCK_GUARD(this->tlsReferenceThread_->tlsLock_);

                AuExchange(this->tlsReferenceThread_->tls_, this->tls_);
                AuExchange(this->tlsReferenceThread_->threadFeatures_, this->threadFeatures_);
                
                bDetachedSuccess = true;
            }
        }

        if (bDetachedSuccess || !bDetached)
        {
            HookReleaseThreadResources();
        }
    }

    // Called whenever an aurora thread is exiting
    void OSThread::HookOnExit()
    {

    }

    // Called whenever a thread object is released, or on thread termination
    // Whichever is soonest
    // It is possible a user may wish to keep a thread handle open
    // Is is expected for a system thread to release on dtor unlike our aurora threads 
    void OSThread::HookReleaseThreadResources()
    {
        tls_.reset();

        for (const auto &feature : threadFeatures_)
        {
            feature->Cleanup();
        }
    }

    void OSThread::AddLastHopeTlsHook(const AuSPtr<AuThreads::IThreadFeature> &feature)
    {
        // TODO: mutex
        if (!feature)
        {
            return;
        }
        AuTryInsert(threadFeatures_, feature);
    }

    void OSThread::Detach()
    {
        detached_ = true;
    }

    AuSPtr<IWaitable> OSThread::AsWaitable()
    {
        return terminated_;
    }

    void OSThread::SendExitSignal()
    {
        exiting_ = true;

        if (this->terminateSignalLs_)
        {
            this->terminateSignalLs_->Set();
        }

        if (this->terminateSignal_)
        {
            this->terminateSignal_->Set();
        }
    }

    bool OSThread::Run()
    {
        if (!terminated_)
        {
            SysPanic("::Run called on system thread");
        }

        if (AuExchange(contextUsed_, true))
        {
            return false;
        }

        terminated_->Reset();

        return ExecuteNewOSContext([=]()
        {
            try
            {
                // this functional backends are being deprecated
                if (info_.callbacks)
                {
                    info_.callbacks->OnEntry(this);
                }
            }
            catch (...)
            {
                Debug::PrintError();
            }
        });
    }

    void OSThread::Exit()
    {
        while ((!terminated_) || (!terminated_->TryLock()))
        {
            if (Exit(false)) break;
        }
    }

    bool OSThread::Exit(bool willReturnToOS)
    {
        if (GetThread() == this)
        {
            if (!this->InternalKill(false))
            {
                return false;
            }

            // release handle + sigterm
            if (!willReturnToOS)
            {
                TeminateOSContext(true);
                while (true)
                {
                    YieldToOtherThread();
                }
            }
        }
        else
        {
            // exit signal 
            exiting_ = true;

            if (!terminated_)
            {
                return true;
            }

            // attempt to join with the thread once it has exited, or timeout 
            if (WaitFor(terminated_.get(), 15 * 1000))
            {
                return true;
            }

            // Do not force terminate if we're marked as dead and still running
            // The thread must've requested suicide and got stuck in a lengthy clean up effort
            if (!exitOnlyOnce_->TryLock())
            {
                AuLogWarn("Watchdog error - OS thread context didn't finish in 15 seconds, but he should exiting now.");
                return false;
            }

            AuLogWarn("Watchdog error - OS thread context didn't finish in 15 seconds, forcefully terminating without a watchdog overlooking onExit");

            // Kill the current OS thread instance
            TeminateOSContext(false);

            // Dispatch kill callback from within an emulated thread context (no switch, just just change thread tls)
            ExecuteInDeadThread([=]()
            {
                this->InternalKill(true);
            });

            exitOnlyOnce_->Unlock();
        }
        return true;
    }

    bool OSThread::Exiting()
    {
        return exiting_;
    }

    void OSThread::SetName(const AuString &name)
    {
        name_ = name;
    }

    EThreadPrio OSThread::GetPrio()
    {
        return prio_;
    }

    AuUInt64 OSThread::GetMask()
    {
        return affinityProcessMask_;
    }

    AuString OSThread::GetName()
    {
        return name_;
    }

    void OSThread::SetAffinity(AuUInt64 mask)
    {
        if (mask == 0) mask = 0xFFFFFFFFFFFFFFFF;
        affinityProcessMask_ = mask;
        UpdateAffinity(mask);
    }

    void OSThread::SetPrio(EThreadPrio prio)
    {
        if (!EThreadPrioIsValid(prio)) return;
        UpdatePrio(prio);
    }

    AuSPtr<TLSView> OSThread::GetTlsView()
    {
        if (!tls_)
        {
            tls_ = AuMakeShared<TLSViewImpl>();
        }

        return tls_;
    }

    bool OSThread::ExecuteNewOSContext(AuFunction<void()> task)
    {
        task_ = task;

        auto ret = SpawnThread([this]()
        {
            this->_ThreadEP();
        }, GetName(), info_.stackSize);
        
        if (ret.first)
        {
            handle_ = (decltype(handle_))ret.second;
        }

        return ret.first;
    }

    void OSThread::_ThreadEP()
    {
        // Poke TLS reference thread entity
        // TODO: we need an internal OSThread *TryPokeTLSThread()
        auto osThread = static_cast<OSThread *>(GetThread());
        this->tlsReferenceThread_ = osThread;

        OSAttach();
        try
        {
            task_();
        }
        catch (...)
        {
            SysPushErrorHAL("OS Thread Aborted");
        }
        Exit(true);
    }

    void OSThread::ExecuteInDeadThread(AuFunction<void()> callback)
    {
        auto old = HandleCurrent();

        if (this == old) [[unlikely]]
        {
            callback();
            return;
        }

        if (old)
        {
            static_cast<OSThread *>(old)->OSDeatach();
        }
        
        this->OSAttach();
        
        callback();

        if (old)
        {
            HandleRegister(old);
            static_cast<OSThread *>(old)->OSDeatach();
        }
        else [[unlikely]]
        {
            HandleRemove();
        }
    }

    void OSThread::UpdateName()
    {
    #if defined(AURORA_IS_MODERNNT_DERIVED)

        if (this->handle_ == INVALID_HANDLE_VALUE)
        {
            return;
        }

        if ((AuBuild::kCurrentPlatform == AuBuild::EPlatform::ePlatformWin32) && (!AuSwInfo::IsWindows10OrGreater()))
        {
            static const DWORD kMSVCExceptionSetName = 0x406D1388;

        #pragma pack(push,8)
            typedef struct tagTHREADNAME_INFO
            {
                DWORD dwType;
                LPCSTR szName;
                DWORD dwThreadID;
                DWORD dwFlags;
            } THREADNAME_INFO;
        #pragma pack(pop)

            if (!IsDebuggerPresent())
            {
                return;
            }

            THREADNAME_INFO info;
            info.dwType = 0x1000;
            info.szName = this->name_.c_str();
            info.dwThreadID = ::GetThreadId(this->handle_);
            info.dwFlags = 0;

            __try
            {
                RaiseException(kMSVCExceptionSetName, 0, sizeof(info) / sizeof(ULONG_PTR), (ULONG_PTR *)&info);
            }
            __except (EXCEPTION_EXECUTE_HANDLER)
            {
            }
        }
        else
        {
            static HRESULT(WINAPI * SetThreadDescription_f)(HANDLE, PCWSTR);
            if (!SetThreadDescription_f)
            {
                SetThreadDescription_f = AuReinterpretCast<decltype(SetThreadDescription_f)>(GetProcAddress(GetModuleHandleW(L"Kernel32.dll"), "SetThreadDescription"));
            }

            if (SetThreadDescription_f)
            {
                SetThreadDescription_f(this->handle_, AuLocale::ConvertFromUTF8(this->name_).c_str());
            }
        }

    #elif defined(AURORA_HAS_PTHREADS)

        pthread_setname_np(handle_, name_.c_str());
    
    #endif
    }

    void OSThread::OSAttach()
    {
        HandleRegister(this);
    #if defined(AURORA_IS_LINUX_DERIVED)
        this->unixThreadId_ = gettid();
    #elif defined(AURORA_IS_XNU_DERIVED)
        this->unixThreadId_ = pthread_getthreadid_np();
    #elif defined(AURORA_IS_BSD_DERIVED)
        #if __FreeBSD_version < 900031
            long lwpid;
            thr_self(&lwpid);
            this->unixThreadId_ = lwpid;
        #elif __FreeBSD_version > 900030
            this->unixThreadId_ = pthread_getthreadid_np();
        #else
            static_assert(false);
        #endif
    #elif defined(AURORA_HAS_PTHREADS)
        this->unixThreadId_ = 0; // !!!!
    #endif
        UpdatePrio(prio_);
        SetAffinity(affinityProcessMask_);
        UpdateName();
    }

    static AuHashMap<EThreadPrio, int> kNiceMap 
    {
        {
            EThreadPrio::ePrioRT, -19
        },
        {
            EThreadPrio::ePrioAboveHigh, -19
        },
        {
            EThreadPrio::ePrioHigh, -11
        },
        {
            EThreadPrio::ePrioNormal, -2
        },
        {
            EThreadPrio::ePrioSub, 5
        }
    };

#if defined(AURORA_IS_MODERNNT_DERIVED)
    static const AuHashMap<EThreadPrio, int> kWin32Map
    {
        {
            EThreadPrio::ePrioRT, THREAD_PRIORITY_TIME_CRITICAL
        },
        {
            EThreadPrio::ePrioAboveHigh, THREAD_PRIORITY_HIGHEST
        },
        {
            EThreadPrio::ePrioHigh, THREAD_PRIORITY_ABOVE_NORMAL
        },
        {
            EThreadPrio::ePrioNormal, THREAD_PRIORITY_NORMAL
        },
        {
            EThreadPrio::ePrioSub, THREAD_PRIORITY_LOWEST
        }
    };
#endif

    void OSThread::UpdatePrio(EThreadPrio prio)
    {
    #if defined(AURORA_IS_MODERNNT_DERIVED)

        if (handle_ == INVALID_HANDLE_VALUE)
        {
            return;
        }

        const int *val;
        if (!AuTryFind(kWin32Map, prio, val))
        {
            return;
        }

        SetThreadPriority(handle_, *val);

    #elif defined(AURORA_HAS_PTHREADS)

        if (!handle_)
        {
            return;
        }

        if (prio == EThreadPrio::ePrioRT)
        {
            sched_param param {};
            param.sched_priority = sched_get_priority_min(SCHED_RR);

            if (pthread_setschedparam(handle_, SCHED_RR, &param) == 0)
            {
                return;
            }

            // fall through on error
        }
        else if (prio_ == EThreadPrio::ePrioRT)
        {
            int policyNonRT = 
            #if defined(AURORA_IS_XNU_DERIVED)
                SCHED_FIFO;
            #else
                SCHED_OTHER;
            #endif  

            sched_param param {};
            param.sched_priority = sched_get_priority_min(policyNonRT);
            pthread_setschedparam(handle_, policyNonRT, &param);
        }

        const int *val;
        if (!AuTryFind(kNiceMap, prio, val))
        {
            return;
        }

        if (auto tid = unixThreadId_)
        {
            setpriority(PRIO_PROCESS, tid, *val);
        }
    #endif

        prio_ = prio;
    }

    void OSThread::UpdateAffinity(AuUInt64 mask)
    {
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        if (handle_ == INVALID_HANDLE_VALUE)
        {
            return;
        }
        SetThreadAffinityMask(handle_, static_cast<DWORD_PTR>(mask));
    #endif
    }

    void OSThread::OSDeatach()
    {

    }

    bool OSThread::InternalKill(bool locked)
    {
        if (terminated_)
        {
            if (!locked)
            {
                if (!exitOnlyOnce_->TryLock())
                {
                    return false;
                }
            }
        }

        HookOnExit();

        try
        {
            // dispatch kill callback
            if (info_.callbacks)
            {
                info_.callbacks->OnExit(this);
            }
        }
        catch (...)
        {
            Debug::PrintError();
            AuLogWarn("Couldn't deinitialize thread");
            AuLogWarn("The smart thing to do at this point would be to panic");
            AuLogWarn("...but we could continue");
            AuLogWarn("Carrying on despite the potential for data integrity loss and memory leaks");
            Telemetry::Mayday();
        }

        HookReleaseThreadResources();

        if (terminated_)
        {
            exitOnlyOnce_->Unlock();
            terminated_->Set();
      
        } 

        if (terminatedSignalLs_)
        {
            terminatedSignalLs_->Set();
        }

        return true;
    }

    void OSThread::TeminateOSContext(bool calledFromThis)
    {
    #if defined(AURORA_IS_MODERNNT_DERIVED)

        if (handle_ == INVALID_HANDLE_VALUE)
        {
            return;
        }

        if (calledFromThis)
        {
            ExitThread(0);
        }
        else
        {
            TerminateThread(handle_, 0);
        }

    #elif defined(AURORA_HAS_PTHREADS)

        pthread_kill(handle_, SIGTERM);

    #else
    
        SysPanic("Not implemented");
    
    #endif
    }

    void OSThread::FreeOSContext()
    {
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        if (auto handle = AuExchange(handle_, {}))
        {
            CloseHandle(handle_); 
        }
    #endif
    }

    AuSPtr<Loop::ILoopSource> OSThread::AsLoopSource()
    {
        return this->terminateSignalLs_;
    }
        
    AuSPtr<IWaitable> OSThread::GetShutdownSignalWaitable()
    {
        return this->terminateSignal_;
    }

    AuSPtr<Loop::ILoopSource> OSThread::GetShutdownSignalLS()
    {
        return this->terminateSignalLs_;
    }

    void InitThreading()
    {
        AuxUlibInitialize();
    }
}