/***
    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::Process
{
    void InitProcessMapLinux()
    {
        //puts("InitProcessMapLinux");
    }

    void DeinitProcessMapLinux()
    {
        //puts("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())
        {
            GetProcFullPath(fileName);
        }

        if (!AuIOFS::FileExists(fileName))
        {
            //printf("drop: %s\n", fileName.c_str());
            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 AuString &line)
        {
            char *endPtr;
            auto base = strtoll(line.c_str(), &endPtr, 16);
            if (errno == ERANGE) return;
            if (*endPtr != '-') return;

            auto end = strtoll(endPtr + 1, &endPtr, 16);
            if (errno == ERANGE) return;
            if (*endPtr != ' ') return;

            auto perms = endPtr + 1;

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

            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;
            }

            if (name.size())
            {
                bool bIsFile = AuIOFS::FileExists(name);
                if (bIsFile)
                {
                    return;
                }
            }

            Section sect {};

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

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

            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;

            AuString 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();
            }

            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 (!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)
            {
                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();
    }
}