#pragma once

#include <shared_mutex>
#include <vector>
#include <Dyld.h>

#include "binaryninjaapi.h"
#include "MachO.h"
#include "VirtualMemory.h"

struct CacheSymbol
{
	BNSymbolType type;
	BNSymbolBinding binding = NoBinding;
	uint64_t address;
	std::string name;

	CacheSymbol() = default;
	CacheSymbol(BNSymbolType type, BNSymbolBinding binding, uint64_t address, std::string name) :
		type(type), binding(binding), address(address), name(std::move(name))
	{}
	~CacheSymbol() = default;

	CacheSymbol(const CacheSymbol& other) = default;
	CacheSymbol& operator=(const CacheSymbol& other) = default;

	CacheSymbol(CacheSymbol&& other) noexcept = default;
	CacheSymbol& operator=(CacheSymbol&& other) noexcept = default;

	std::pair<std::string, BinaryNinja::Ref<BinaryNinja::Type>> DemangledName(BinaryNinja::BinaryView& view) const;

	// NOTE: you should really only call this when adding the symbol to the view.
	std::pair<BinaryNinja::Ref<BinaryNinja::Symbol>, BinaryNinja::Ref<BinaryNinja::Type>> GetBNSymbolAndType(BinaryNinja::BinaryView& view) const;
};

enum class CacheRegionType
{
	Image,
	StubIsland,
	DyldData,
	NonImage,
};

struct CacheRegion
{
	CacheRegionType type;
	std::string name;
	uint64_t start;
	uint64_t size;
	// Associate this region with this image, this makes it easier to identify what image owns this region.
	std::optional<uint64_t> imageStart;
	BNSegmentFlag flags;

	CacheRegion() = default;
	~CacheRegion() = default;

	CacheRegion(const CacheRegion& other) = default;
	CacheRegion& operator=(const CacheRegion& other) = default;

	CacheRegion(CacheRegion&& other) noexcept = default;
	CacheRegion& operator=(CacheRegion&& other) noexcept = default;

	AddressRange AsAddressRange() const { return {start, start + size}; }

	BNSectionSemantics SectionSemanticsForRegion() const
	{
		if ((flags & SegmentExecutable) && (flags & SegmentDenyWrite))
			return ReadOnlyCodeSectionSemantics;

		if (flags & SegmentExecutable)
			return DefaultSectionSemantics;

		if (flags & SegmentDenyWrite)
			return ReadOnlyDataSectionSemantics;

		return ReadWriteDataSectionSemantics;
	}
};

// Represents a single image and its associated memory regions.
struct CacheImage
{
	uint64_t headerAddress;
	std::string path;
	// A list to the start of memory regions associated with the image.
	// This lets us load all regions for a given image easily.
	std::vector<uint64_t> regionStarts;
	std::shared_ptr<SharedCacheMachOHeader> header;

	CacheImage() = default;
	~CacheImage() = default;

	CacheImage(const CacheImage& other) = default;
	CacheImage& operator=(const CacheImage& other) = default;

	CacheImage(CacheImage&& other) noexcept = default;
	CacheImage& operator=(CacheImage&& other) noexcept = default;

	// Get the file name from the path.
	std::string GetName() const { return BaseFileName(path); }

	// Get the names of the dependencies.
	std::vector<std::string> GetDependencies() const;
};

enum class CacheEntryType
{
	Primary,
	Secondary,
	// A special entry that holds symbols for other cache entries.
	Symbols,
	// If the type is marked as this then all mappings will be marked as such.
	DyldData,
	// A single stub mapping file.
	Stub,
};

// Describes a single files cache information
class CacheEntry
{
	CacheEntryType m_type;
	std::string m_filePath;
	std::string m_fileName;
	dyld_cache_header m_header {};
	// Mappings tell us _where_ to map the regions within the flat address space.
	// Without this we wouldn't know where the entry is supposed to exist in the address space.
	std::vector<dyld_cache_mapping_info> m_mappings {};
	// Mapping of image path to image info, used within ProcessImagesAndRegions to add them to the cache.
	// Also used to retrieve the image dependencies.
	std::vector<std::pair<std::string, dyld_cache_image_info>> m_images {};
	std::shared_ptr<MappedFileRegion> m_file;

public:
	CacheEntry(std::string filePath, std::string fileName, CacheEntryType type, dyld_cache_header header,
		std::vector<dyld_cache_mapping_info> mappings, std::vector<std::pair<std::string, dyld_cache_image_info>> images,
		std::shared_ptr<MappedFileRegion> file);

	CacheEntry() = default;
	CacheEntry(const CacheEntry&) = default;

	CacheEntry(CacheEntry&&) = default;
	CacheEntry& operator=(CacheEntry&&) = default;

	// Construct a cache entry from the file on disk.
	static CacheEntry FromFile(const std::string& filePath, const std::string& fileName, CacheEntryType type);

	const std::shared_ptr<MappedFileRegion>& GetFile() const { return m_file; }

	// Get the headers virtual address.
	// This is useful if you need to read relative to the start of the entry file.
	std::optional<uint64_t> GetHeaderAddress() const;

	// Get the mapped address for a given file offset.
	// Ex. passing 0x0 will retrieve the mapped address for the start of the file (i.e. the header)
	std::optional<uint64_t> GetMappedAddress(uint64_t fileOffset) const;

	CacheEntryType GetType() const { return m_type; }
	// Ex. "/myuser/mypath/dyld_shared_cache_arm64e"
	const std::string& GetFilePath() const { return m_filePath; }
	// Ex. "dyld_shared_cache_arm64e"
	const std::string GetFileName() const { return m_fileName; }
	const dyld_cache_header& GetHeader() const { return m_header; }
	const std::vector<dyld_cache_mapping_info>& GetMappings() const { return m_mappings; }
	const std::vector<std::pair<std::string, dyld_cache_image_info>>& GetImages() const { return m_images; }
};

// The ID for a given CacheEntry, use this instead of passing a pointer around to avoid complexity :V
typedef uint32_t CacheEntryId;

// The C in DSC.
// This represents the entire cache, all regions and images are visible from here.
// This is the dump for all the information, and what the workflow activities and the UI want.
// Creating this is expensive, both in actual processing and just copying, so we only generate this
// once every time the database is open.
class SharedCache
{
	// Calculated within `AddEntry`, this indicates where the shared cache image is based at.
	uint64_t m_baseAddress = 0;
	// The shared cache can own the virtual memory, this is fine...
	std::shared_ptr<VirtualMemory> m_vm;
	std::vector<CacheEntry> m_entries {};
	// This information is used in tandem with the cache images to load memory regions into the binary view.
	AddressRangeMap<CacheRegion> m_regions {};
	// Describes the images of the cache.
	std::unordered_map<uint64_t, CacheImage> m_images {};
	// All the external symbols for this cache. Both mapped and unmapped (not in the view).
	std::unordered_map<uint64_t, CacheSymbol> m_symbols {};
	// Quickly lookup a symbol by name, populated by `FinalizeSymbols`.
	// `m_namedSymbols` is modified in a worker thread spawned by view init so we must not get a symbol until its populated.
	std::unordered_map<std::string, uint64_t> m_namedSymbols {};
	// Used to guard `m_namedSymbols` as it's accessed on multiple threads.
	// NOTE: Wrapped in unique_ptr to keep SharedCache movable.
	std::unique_ptr<std::shared_mutex> m_namedSymMutex;

	// Local symbols entry and its mapping, used to read symbol tables from the .symbols file.
	// These are handled separately as they are not mapped into the main virtual memory of the cache.
	std::optional<CacheEntry> m_localSymbolsEntry;
	std::shared_ptr<VirtualMemory> m_localSymbolsVM;

	bool ProcessEntryImage(const std::string& path, const dyld_cache_image_info& info);

	// Add a region known not to overlap with another, otherwise use AddRegion.
	// returns whether the region was inserted.
	bool AddNonOverlappingRegion(CacheRegion region);

public:
	SharedCache() = default;
	explicit SharedCache(uint64_t addressSize);

	SharedCache(const SharedCache &) = delete;
	SharedCache &operator=(const SharedCache &) = delete;

	SharedCache(SharedCache &&) noexcept = default;
	SharedCache &operator=(SharedCache &&) noexcept = default;

	uint64_t GetBaseAddress() const { return m_baseAddress; }
	std::shared_ptr<VirtualMemory> GetVirtualMemory() const { return m_vm; }
	const std::vector<CacheEntry>& GetEntries() const { return m_entries; }
	const AddressRangeMap<CacheRegion>& GetRegions() const { return m_regions; }
	const std::unordered_map<uint64_t, CacheImage>& GetImages() const { return m_images; }
	const std::unordered_map<uint64_t, CacheSymbol>& GetSymbols() const { return m_symbols; }

	const std::optional<CacheEntry>& GetLocalSymbolsEntry() const { return m_localSymbolsEntry; }
	std::shared_ptr<VirtualMemory> GetLocalSymbolsVM() const { return m_localSymbolsVM; }

	void AddImage(CacheImage&& image);

	// Add a region that may overlap with another.
	void AddRegion(CacheRegion&& region);

	void AddSymbol(CacheSymbol symbol);

	void AddSymbols(std::vector<CacheSymbol>&& symbols);

	// Adds the cache entry and populates the virtual memory using the mapping information.
	// After being added the entry is read only, there is nothing that can modify it.
	void AddEntry(CacheEntry entry);

	void ProcessEntryImages(const CacheEntry& entry);

	void ProcessEntryRegions(const CacheEntry& entry);

	void ProcessEntrySlideInfo(const CacheEntry& entry) const;

	// Construct the named symbols lookup map for use with `GetSymbolWithName`.
	void ProcessSymbols();

	std::optional<CacheEntry> GetEntryContaining(uint64_t address) const;

	std::optional<CacheEntry> GetEntryWithImage(const CacheImage& image) const;

	std::optional<CacheRegion> GetRegionAt(uint64_t address) const;

	std::optional<CacheRegion> GetRegionContaining(uint64_t address) const;

	std::optional<CacheImage> GetImageAt(uint64_t address) const;

	std::optional<CacheImage> GetImageContaining(uint64_t address) const;

	// TODO: Rename to GetImageWithPath and then make another one for the image name.
	std::optional<CacheImage> GetImageWithName(const std::string& name) const;

	std::optional<CacheSymbol> GetSymbolAt(uint64_t address) const;

	std::optional<CacheSymbol> GetSymbolWithName(const std::string& name);
};