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

    File: Process.Win32.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <RuntimeInternal.hpp>
#include "Processes.hpp"
#include "Process.Win32.hpp"
#include <shellapi.h>
#include <tlhelp32.h>

namespace Aurora::Processes
{
    static HANDLE gLeaderJob;

    void InitWin32()
    {
        gLeaderJob = CreateJobObject(NULL, NULL); 
        if (!gLeaderJob)
        {
            SysPushErrorArg("CreateJobObject error");
            return;
        }

        JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli {};
        jeli.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
        if (!SetInformationJobObject(gLeaderJob, JobObjectExtendedLimitInformation, &jeli, sizeof(jeli)))
        {
            SysPushErrorGen("SetInformationJobObject error");
        }
    }

    class ProcessImpl : public IProcess
    {
    public:
        ProcessImpl(AuString execModule, AuList<AuString> args);
        ~ProcessImpl();

        bool TermWinEnumProcesses();

        bool TryKill() override;
        bool HasExited();

        bool Terminate() override;
        Aurora::Threading::IWaitable *AsWaitable() override;
        AuSInt GetExitCode() override;

        void ShutdownPipes();
        
        bool Read(bool error, void *buffer, AuUInt32 &len) override;
        bool Write(const void *buffer, AuUInt32 len) override;

        bool Start(enum ESpawnType type, bool fwdOut, bool fwdErr, bool fwdIn) override;

    private:
        
        HANDLE pipeStdOutRead_  {};
        HANDLE pipeStdOutWrite_ {};
        HANDLE pipeStdErrRead_  {};
        HANDLE pipeStdErrWrite_ {};
        HANDLE pipeStdInRead_   {};
        HANDLE pipeStdInWrite_  {};
        
        AuString execModule_;
        ESpawnType type_;
        
        AuList<AuString> args_;
        AuList<const char *> cargs_;
        AuString windowsCli_;
        
        Threading::Threads::ThreadUnique_t thread_;
        HANDLE process_;
        HANDLE hthread_;
        AuSInt exitCode_;
    };


    ProcessImpl::ProcessImpl(AuString execModule, AuList<AuString> args) : execModule_(execModule), args_(args)
    {
        this->args_.insert(this->args_.begin(), execModule);
               
        // ehhhh https://github.com/tritao/WindowsSDK/blob/07983c7ba4f6861d15e23f195744c60c0c249ce0/SDKs/SourceDir/Windows%20Kits/10/Source/10.0.17763.0/ucrt/exec/cenvarg.cpp#L23
        for (const auto &arg : this->args_)
        {
            this->cargs_.push_back(arg.c_str());
            this->windowsCli_ += arg + " ";
        }

        this->cargs_.push_back(nullptr);
        this->windowsCli_.resize(this->windowsCli_.size() - 1);
    }

    ProcessImpl::~ProcessImpl()
    {
        if (this->type_ == ESpawnType::eSpawnChildProcessWorker)
        {
            TryKill();
            Terminate();
        }

        if (this->thread_)
        {
            this->thread_.reset();
        }

        CloseHandle(this->process_);
        CloseHandle(this->hthread_);
  
        ShutdownPipes();
    }

    static BOOL TermWinHandleWin32Thread(HWND handle, LPARAM a)
    {
        SendMessageA(handle, WM_CLOSE, 0, 0);
        (*reinterpret_cast<AuUInt32 *>(a))++;
        return true;
    }

    bool ProcessImpl::TermWinEnumProcesses()
    {
        THREADENTRY32 te{};
        HANDLE h{};

        if ((!this->process_) || (this->process_ == INVALID_HANDLE_VALUE))
        {
            return false;
        }

        h = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, GetProcessId(this->process_));
        if (h == INVALID_HANDLE_VALUE)
        {
            return false;
        }

        te.dwSize = sizeof(te);
        AuUInt32 count{};
        if (Thread32First(h, &te))
        {
            do
            {
                if (te.dwSize >= FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) +
                    sizeof(te.th32OwnerProcessID))
                {
                    EnumThreadWindows(te.th32ThreadID, TermWinHandleWin32Thread, reinterpret_cast<LPARAM>(&count));
                }
                te.dwSize = sizeof(te);
            } while (Thread32Next(h, &te));
        }

        CloseHandle(h);

        if (!count)
        {
            return false;
        }
            
        return Threading::YieldPoll(true, 2500, [=]()
        {
            return !HasExited();
        });
    }

    bool ProcessImpl::TryKill()
    {
        return TermWinEnumProcesses();
    }

    bool ProcessImpl::HasExited()
    {
        DWORD a;
        if (!GetExitCodeProcess(this->process_, &a)) return true;
        return  a != STILL_ACTIVE;
    }

    bool ProcessImpl::Terminate()
    {
        if ((this->process_) && (this->process_ != INVALID_HANDLE_VALUE))
        {
            return TerminateProcess(this->process_, 0) || HasExited();
        }
        return false;
    }

    Aurora::Threading::IWaitable *ProcessImpl::AsWaitable()
    {
        if (!this->thread_) return nullptr;
        return this->thread_->AsWaitable();
    }

    AuSInt ProcessImpl::GetExitCode()
    {
        return this->exitCode_;
    }

    void ProcessImpl::ShutdownPipes()
    {
        if (auto handle = std::exchange(this->pipeStdOutRead_, {}))
            CloseHandle(handle);

        if (auto handle = std::exchange(this->pipeStdOutWrite_, {}))
            CloseHandle(handle);

        if (auto handle = std::exchange(this->pipeStdErrRead_, {}))
            CloseHandle(handle);

        if (auto handle = std::exchange(this->pipeStdErrWrite_, {}))
            CloseHandle(handle);

        if (auto handle = std::exchange(this->pipeStdInRead_, {}))
            CloseHandle(handle);

        if (auto handle = std::exchange(this->pipeStdInWrite_, {}))
            CloseHandle(handle);
    }
        
    bool ProcessImpl::Read(bool error, void *buffer, AuUInt32 &len)
    {
        DWORD size = std::exchange(len, 0);
        auto handle = error ? pipeStdErrRead_ : pipeStdOutRead_;
        if (!handle) return false;
        auto ret = ReadFile(handle, buffer, size, &size, NULL);
        len = size;
        return ret;
    }

    bool ProcessImpl::Write(const void *buffer, AuUInt32 len)
    {
        DWORD size = len;
        if (!pipeStdInWrite_) return false;
        return WriteFile(pipeStdInWrite_, buffer, size, &size, NULL) && size == len;
    }

    bool ProcessImpl::Start(enum ESpawnType type, bool fwdOut, bool fwdErr, bool fwdIn)
    {
        this->exitCode_ = 0x10110100;
        this->type_ = type;
            
        if (type == ESpawnType::eSpawnAtomicOvermap)
        {
            _spawnv(_P_OVERLAY, this->execModule_.c_str(), this->cargs_.data());
            SysPushErrorGen("_spawnv didn't overwrite the process map, given {} ({})", this->execModule_, this->windowsCli_);
            return false;
        }

        {
            SECURITY_ATTRIBUTES saAttr{};

            saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
            saAttr.bInheritHandle = TRUE;
            saAttr.lpSecurityDescriptor = NULL;

            if (fwdOut)
            {
                if (!CreatePipe(&pipeStdOutRead_, &pipeStdOutWrite_, &saAttr, 0))
                {
                    return false;
                }

                if (!SetHandleInformation(pipeStdOutRead_, HANDLE_FLAG_INHERIT, 0))
                {
                    return false;
                }
            }
                
            if (fwdErr)
            {
                if (!CreatePipe(&pipeStdErrRead_, &pipeStdErrWrite_, &saAttr, 0))
                {
                    return false;
                }

                if (!SetHandleInformation(pipeStdErrRead_, HANDLE_FLAG_INHERIT, 0))
                {
                    return false;
                }
            }

            if (fwdIn)
            {
                if (!CreatePipe(&pipeStdInRead_, &pipeStdInWrite_, &saAttr, 0))
                {
                    return false;
                }

                if (!SetHandleInformation(pipeStdInWrite_, HANDLE_FLAG_INHERIT, 0))
                {
                    return false;
                }
            }
        }

        Threading::Threads::AbstractThreadVectors handler;

        PROCESS_INFORMATION processInfo = { 0 };
        {
            STARTUPINFOW startupInfo = { 0 };
            startupInfo.cb = sizeof(startupInfo);
                
            startupInfo.hStdInput  = pipeStdInRead_;
            startupInfo.hStdError  = pipeStdErrWrite_;
            startupInfo.hStdOutput = pipeStdOutWrite_;
            startupInfo.dwFlags   |= ((fwdIn || fwdErr || fwdOut) ? STARTF_USESTDHANDLES : 0);

            auto result = CreateProcessW(Locale::ConvertFromUTF8(this->execModule_).c_str(),
                                         Locale::ConvertFromUTF8(this->windowsCli_).data(),
                                         NULL, NULL, FALSE,
                                         NULL,
                                         NULL, NULL, &startupInfo, &processInfo);

            if (!result)
            {
                SysPushErrorGen("CreateProcess failed");
                return false;
            }

            this->process_ = processInfo.hProcess;
            this->hthread_ = processInfo.hThread;

            if (type == ESpawnType::eSpawnChildProcessWorker)
            {
                if (gLeaderJob)
                {
                    if (!AssignProcessToJobObject(gLeaderJob, processInfo.hProcess))
                    {
                        SysPushErrorGen("Could not AssignProcessToObject");
                    }
                }
            }
        }

        handler.DoRun = [=](Threading::Threads::IAuroraThread *)
        {
            {
                WaitForSingleObject(processInfo.hProcess, INFINITE);

                DWORD exitCode;
                auto result = GetExitCodeProcess(processInfo.hProcess, &exitCode);
                this->exitCode_ = exitCode;
            }
        };

        this->thread_ = Threading::Threads::ThreadUnique(handler);
        if (!this->thread_) return false;

        this->thread_->Run();
        return true;
    }

    AUKN_SYM IProcess *SpawnNew(const AuString &app, const AuList<AuString> &args)
    {
        return _new ProcessImpl(app, args);
    }

    AUKN_SYM void SpawnRelease(IProcess *process)
    {
        SafeDelete<ProcessImpl *>(process);
    }
}