#pragma once
#include "types.hpp"
#include "loader/binary.hpp"
#include "disasm/engine.hpp"
#include <map>
#include <set>
#include <vector>

class CFGBuilder
{
    public:
        CFGBuilder(Binary* binary, Disassembler* disasm, bool subsOnly = false);
        CFG build();
        static void clean(CFG& cfg);

    private:
        void buildGotMap();
        Function buildFunction(addr_t start, const string& name);
        BasicBlock disassembleBlock(addr_t start, addr_t& next);
        set<addr_t> findPrologs();
        void scanDataPointers();
        addr_t skipEndbr64(addr_t addr);
        bool isPLTStub(const Function& func) const;
        void markVisited(addr_t addr);
        bool isVisited(addr_t addr);

        addr_t resolveGOT(const Instruction& inst);
        addr_t resolveJumpTable(const vector<Instruction>& block, const Instruction& inst);
        addr_t traceReg(const string& reg, const vector<Instruction>& block, size_t idx);
        void recordIndirect(const Instruction& inst, addr_t target, bool ok, const string& name, const string& lib);

        Binary* m_bin;
        Disassembler* m_dis;
        bool m_subs_only;
        set<addr_t> m_funcs;
        set<addr_t> m_blocks;
        set<addr_t> m_prologs;
        map<addr_t, pair<string, string>> m_got;
        vector<IndirectTarget> m_targets;
};
