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

    File: ModuleInfo.Linux.cpp
    Date: 2022-4-20
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuProcessMap.Linux.hpp"
#include "AuProcessMap.hpp"
#include <link.h>

#if defined(AURORA_IS_64BIT)
    using Elf_Ehdr = Elf64_Ehdr;
    using Elf_Shdr = Elf64_Shdr;
#else
    using Elf_Ehdr = Elf32_Ehdr;
    using Elf_Shdr = Elf32_Shdr;
#endif

namespace Aurora::CmdLine
{
    extern AuString gHackLoader;
}

namespace Aurora::Process
{
    void ForceReplacePathHack(const AuString &path);

    void InitProcessMapLinux()
    {

    }

    void DeinitProcessMapLinux()
    {

    }

    static AuHashMap<AuUInt, AuTuple<AuString, AuList<AuTuple<AuUInt, AuUInt, AuUInt, AuUInt>>>> gModuleMap;

    static int LdDlIIterateCallback(struct dl_phdr_info *info,
                                    size_t size,
                                    void *data)
    {
        AuString fileName = info->dlpi_name[0] ? info->dlpi_name : AuString{};
        if (fileName.empty())
        {
            if (auto optProcessPath = GetProcessFullPath())
            {
                fileName = *optProcessPath;
            }
        }
        else
        {
            if (!AuIOFS::FileExists(fileName))
            {
                // NO FILE linux-vdso.so.1
                return 0;
            }
        }

        AuList<AuTuple<AuUInt, AuUInt, AuUInt, AuUInt>> sections;
        for (int j = 0; j < info->dlpi_phnum; j++)
        {
            AuTryInsert(sections, AuMakeTuple((AuUInt)info->dlpi_phdr[j].p_vaddr, (AuUInt)info->dlpi_phdr[j].p_offset, (AuUInt)info->dlpi_phdr[j].p_memsz, (AuUInt)info->dlpi_phdr[j].p_flags));
        }

        AuTryInsert(gModuleMap, info->dlpi_addr, AuMakeTuple(fileName, sections));
        
        return 0;
    }

    static void PassOneScanLd()
    {
        ::dl_iterate_phdr(LdDlIIterateCallback, NULL);
    }

    static void PassTwoScanMaps()
    {
        AuString map;

        if (!AuIOFS::ReadString("/proc/self/maps", map))
        {
            return;
        }

        Sections sections;
        AuParse::SplitNewlines(map, [&](const AuROString &line)
        {
            bool bIsSpecialFile {};
            const char *endPtr = line.data() + line.size();
            const char *endPtr2 = line.data() + line.size();
            const char *endPtr3 = line.data() + line.size();
            auto optBase = AuParse::ParseUInt16(line.data(), endPtr);
            if (!optBase) return;
            auto base = optBase.value();
            if (*endPtr != '-') return;

            auto optEnd = AuParse::ParseUInt16(endPtr + 1, endPtr2);
            if (!optEnd) return;
            auto end = optEnd.value();
            if (*endPtr2 != ' ') return;

            auto perms = endPtr2 + 1;

            auto A = line.find_first_of(':');
            if (A == AuString::npos) return;

            auto offsetStart = perms + 5;
            auto optuFileOffset = AuParse::ParseUInt16(offsetStart, endPtr3);
            AuUInt uFileOffset = optuFileOffset ? optuFileOffset.value() : 0u;

            A += 4;
            auto gross = line.substr(A);

            AuString name;
            auto B = gross.find_first_of(' ');
            if (B != AuString::npos) 
            {
                auto C = gross.find_first_not_of(' ', B);
                if (C != AuString::npos)
                {
                    name = gross.substr(C);
                }
            }

            if (name == "[stack]")
            {
                name = kSectionNameStack;
            }
            else if (name == "[heap]")
            {
                name = kSectionNameHeap;
            }
            else if (name == "[vdso]")
            {
                name = kSectionNameFile;
            }

            AuOptional<AuROString> optProcessPath;
            if (AuCmdLine::gHackLoader.size())
            {
                optProcessPath = GetProcessFullPath();
            }

            if (name.size())
            {
                bIsSpecialFile = AuIOFS::FileExists(name);

                if (bIsSpecialFile)
                {
                    for (auto &[uBaseAddress, pathSectionPair] : gModuleMap)
                    {
                        if (uBaseAddress == base)
                        {
                            if (optProcessPath)
                            {
                                if (AuGet<0>(pathSectionPair) == optProcessPath.Value())
                                {
                                    ForceReplacePathHack(name);
                                    AuGet<0>(pathSectionPair) = name;
                                    AuResetMember(optProcessPath);
                                }
                            }

                            return;
                        }

                        for (auto &[a, b, c, d] : AuGet<1>(pathSectionPair))
                        {
                            if ((uBaseAddress + a) <= base &&
                                (uBaseAddress + a + c) > base)
                            {
                                return;
                            }
                        }
                    }
                }
            }

            Section sect {};

            sect.origVa      = 0;
            sect.baseVa      = base;
            sect.size        = end - base;
            sect.fsOff       = uFileOffset;

            sect.pt.readable = perms[0] == 'r';
            sect.pt.writable = perms[1] == 'w';
            sect.pt.NX       = perms[2] != 'x';

            if (bIsSpecialFile)
            {
                // edge case: these are like windows nt section views
                // insert these as their own modules
                // (we cant present Sections of weakly owned files )
                static AuMutex gLock;
                static AuHashMap<AuString, AuSPtr<PublicModule>> gMap;

                sect.name = kSectionNameFile;

                AU_LOCK_GUARD(gLock);

                auto &refMod = gMap[name];
                if (refMod)
                {
                    // already cached
                } 
                else
                {
                    refMod = AuMakeShared<PublicModule>();
                    if (!refMod) return;

                    refMod->moduleMeta = AuMakeShared<ModuleMeta>();
                    if (!refMod->moduleMeta) return;

                    AuROString fileName;
                    AuIOFS::GetFileFromPath(fileName, name);
                    refMod->moduleMeta->moduleBase = base;
                    refMod->moduleMeta->moduleName = fileName;
                    refMod->moduleMeta->modulePath = name;
                    refMod->moduleMeta->origVa = 0;
                }

                sect.moduleMeta = refMod;
            }
            else
            {
                // most likely
                sect.name = AuMove(name);
            }

            AuTryInsert(sections, sect);
        }, false);

        BorrowOtherSectionArray([&](AuList<Section> &out)
        {
            out.clear();
            out.insert(out.end(), sections.begin(), sections.end());
        });
    }

    static void PassThreeGatherMissingProgNames()
    {
        for (const auto & [baseAddress, pair] : gModuleMap)
        {
            auto & [path, sections] = pair;

            AuROString file;
            AuIOFS::GetFileFromPath(file, path);
            auto object = AuIOFS::OpenReadUnique(path);
            if (!object)
            {
                SysPushErrorIO();
                continue;
            }

            Elf_Ehdr header;
            AuUInt read;
            if (!object->Read(AuMemoryViewStreamWrite(AuMemoryViewWrite(&header, sizeof(header)), read)))
            {
                SysPushErrorIO();
                continue;
            }

            AuList<Elf_Shdr> elfSections;
            if (!AuTryResize(elfSections, header.e_shnum))
            {
                SysPushErrorMem();
                continue;
            }

            object->SetOffset(header.e_shoff);
            if (!object->Read(AuMemoryViewStreamWrite(AuMemoryViewWrite(elfSections.data(),
                                                                        elfSections.size() * sizeof(Elf_Shdr)
                                                                        ),
                                                      read)))
            {
                SysPushErrorIO();
                continue;
            }

            auto &strtab = elfSections[header.e_shstrndx];

            AuList<AuUInt8> strHeap;
            if (strtab.sh_size > 80 * 1024 * 1024)
            {
                SysPushErrorMem();
                continue;
            }

            if (!AuTryResize(strHeap, strtab.sh_size))
            {
                SysPushErrorMem();
                continue;
            }

            object->SetOffset(strtab.sh_offset);
            if (!object->Read(AuMemoryViewStreamWrite(AuMemoryViewWrite(strHeap.data(),
                                                                        strHeap.size()),
                                                      read)))
            {
                SysPushErrorIO();
                continue;
            }

            Sections modSections;
            for (const auto & section : elfSections)
            {
                if (section.sh_name > strHeap.size())
                {
                    continue;
                }
                
                auto len = strnlen((const char *)strHeap.data() + section.sh_name, strHeap.size() - section.sh_name);
                AuString sectionName((const char *)strHeap.data() + section.sh_name, len);
                AuUInt fileOffset = section.sh_offset;
                AuUInt programOffset = section.sh_addr;

                Section sect {};
                sect.origVa = section.sh_addr;
                sect.baseVa = 0;

                for (const auto & [base, offset, size, flags] : sections)
                {
                    if ((offset <= fileOffset) &&
                       ((offset + size) > fileOffset) && fileOffset && offset)
                    {
                        sect.pt.readable = section.sh_addr ? (section.sh_flags & SHF_ALLOC) : false;//?

                        auto baseOfMap = baseAddress;
                        auto offsetInMap = fileOffset - offset;

                        sect.baseVa = (baseOfMap + base) + offsetInMap;
                    }
                }

                sect.size   = section.sh_size;
                sect.fsOff  = fileOffset;

                sect.pt.NX       = !(section.sh_flags & SHF_EXECINSTR);
                sect.pt.writable = section.sh_flags & SHF_WRITE;

                if (!sect.pt.readable)
                {
                    sect.baseVa = 0;
                }

                sect.name = sectionName;

                if (sect.fsOff || sect.size)
                {
                    AuTryInsert(modSections, sect);
                }
            }

            auto pub = AuMakeShared<PublicModule>();
            if (!pub) return;

            pub->moduleMeta = AuMakeShared<ModuleMeta>();
            if (!pub->moduleMeta) return;

            pub->moduleMeta->moduleBase = baseAddress;
            pub->moduleMeta->moduleName = AuMove(file);
            pub->moduleMeta->modulePath = AuMove(path);
            pub->moduleMeta->origVa = 0;

            pub->sections = AuMove(modSections);

            // TODO: ...
            InsertModuleCache(ModuleBasePair{pub->moduleMeta->moduleName, baseAddress + AuGet<0>(sections[0])}, pub);
        }
    }

    static void PassFourAdd()
    {
        // TODO: ?
    }

    void RescanMaps()
    {
        PassOneScanLd();
        PassTwoScanMaps();
        PassThreeGatherMissingProgNames();
        PassFourAdd();
    }
}