// SPDX-License-Identifier: MIT
#pragma once
#include "Interface/Core/JIT/Relocations.h"
#include "Interface/Core/Frontend.h"
#include "Interface/Core/CPUBackend.h"

#include <FEXCore/Core/CodeCache.h>
#include <FEXCore/Core/Context.h>
#include <FEXCore/Config/Config.h>
#include <FEXCore/Utils/File.h>
#include <FEXCore/Utils/WorkQueueThread.h>
#include <FEXCore/fextl/memory.h>
#include <FEXCore/fextl/string.h>
#include <FEXCore/fextl/unordered_set.h>
#include <FEXCore/fextl/robin_map.h>
#include <FEXCore/fextl/map.h>
#include <FEXCore/fextl/vector.h>
#include <stdint.h>
#include <mutex>
#include <optional>
#include <span>
#include <xxhash.h>

namespace FEXCore {

namespace Context {
  class ContextImpl;
}

namespace DiskCache {

  namespace MesaFOZ {

#define FOSSILIZE_BLOB_HASH_LENGTH 40 /* SHA1 hexadecimal string length */

    struct __attribute__((packed)) foz_payload_key {
      uint8_t bytes[FOSSILIZE_BLOB_HASH_LENGTH];
    };

    struct __attribute__((packed)) foz_payload_header {
      uint32_t payload_size;
      uint32_t format;
      uint32_t crc;
      uint32_t uncompressed_size;
    };

  } // namespace MesaFOZ

  class IndexedDB;

  struct IndexEntry {
    IndexedDB* DB;
    uint64_t Offset;
    uint32_t Size;
    uint32_t GuestSize;
    XXH128_hash_t GuestHash;
    fextl::vector<uint32_t> GuestExtents;
  };

  struct IndexCacheHead {
    struct IndexEntry MainEntry;
    uint64_t MainEntryFootprint;
    fextl::unique_ptr<fextl::multimap<uint64_t, IndexEntry>> MoreEntries; // sorted by guest footprint
  };

  struct __attribute__((packed)) BlobFixedHeader {
    uint32_t GuestSize;
    uint32_t HostSize;
    uint32_t EntryPointCount;
    uint32_t SmallRelocCount;
    uint32_t ThunkRelocCount;
    XXH128_hash_t GuestHash;
  };

  // packed struct for types 0, 2 and 3. type 1 is bigger and separate below
  struct __attribute__((packed)) BlobSmallRelocation {
    uint32_t Offset;
    uint8_t Type;
    union {
      struct __attribute__((packed)) {
        uint32_t Symbol;
      } Named;
      struct __attribute__((packed)) {
        uint64_t GuestRIP;
      } RIPLiteral;
      struct __attribute__((packed)) {
        uint8_t RegisterIndex;
        uint64_t GuestRIP;
      } RIPMove;
      struct __attribute__((packed)) {
        uint8_t RegisterIndex;
        uint8_t ValueSize;
        uint32_t SiteOffset;
      } PatchableData;
    };
  };

  // type 1, implicit
  struct __attribute__((packed)) BlobThunkRelocation {
    uint32_t Offset;
    uint8_t RegisterIndex;
    uint8_t SymbolHash[32]; // sha256sum in the real RelocNamedThunkMove
  };

  struct CodeHitData {
    fextl::vector<uint8_t> Blob;
    std::span<uint8_t> HostCode;
    std::span<const uint64_t> GuestPages;
    std::span<uint64_t> EntryPointRIPs;
    std::span<const uint32_t> EntryPointHostOffsets;

    // the spans above point to memory owned by the Blob vec, so it's important this can't be copied
    CodeHitData() = default;
    CodeHitData(CodeHitData&&) = default;
    CodeHitData& operator=(CodeHitData&&) = default;
    CodeHitData(const CodeHitData&) = delete;
    CodeHitData& operator=(const CodeHitData&) = delete;
  };

  using Index = fextl::robin_map<uint64_t, IndexCacheHead>;

  class FOZFile {
  public:
    bool Open(const fextl::string& CacheFileName, bool ReadOnly);
    bool Lock(uint32_t TimeoutMS) {
      if (!FD) {
        return false;
      }
      return FD->Lock(TimeoutMS);
    }
    bool Unlock() {
      if (!FD) {
        return false;
      }
      return FD->Unlock();
    }
    File::File::FileHandleType GetHandle() {
      return FD ? FD->GetHandle() : (File::File::FileHandleType)-1;
    }
    ssize_t Size();
    bool ReadAll(fextl::vector<uint8_t>& Out); // from first blob
    bool ReadBlob(uint64_t Offset, std::span<uint8_t> OutBlob);
    bool WriteBlob(const MesaFOZ::foz_payload_key& Key, std::span<const std::span<const uint8_t>> BlobChunks, uint64_t& OutBlobOffset);

  private:
    static constexpr uint32_t OPEN_LOCK_TIMEOUT_MS = 100;

    fextl::string FileName;
    fextl::unique_ptr<File::File> FD;
    bool ReadOnly = false;
  };

  class IndexedDB {
  public:
    bool Open(const fextl::string& CacheDBName, bool ReadOnly);
    void PopulateIndex(Index& CacheIndex, bool& FoundMetadata);
    bool ReadCacheBlob(uint64_t Offset, std::span<uint8_t> OutBlob);
    bool StoreCacheBlob(const MesaFOZ::foz_payload_key& UniqueKey, uint64_t LookupKey, std::span<const uint8_t> Blob, Index& CacheIndex,
                        std::mutex& IndexMutex, std::span<const uint8_t> IndexBlob);

  private:
    // stores run on the Writer, so returning quick isn't as important
    static constexpr uint32_t STORE_LOCK_TIMEOUT_MS = 1000;
    static constexpr uint64_t BIG_MAPPING_SIZE = 1ULL << 33;
    static constexpr uint32_t LOOKUP_KEY_MAX_BUCKET_DEPTH = 20;

    FOZFile CacheFOZ;
    uint8_t* CacheFileMapping = nullptr;
    std::atomic<uint64_t> CacheFileSize;
    FOZFile IndexFOZ;
    bool ReadOnly = false;
  };

  class DiskCache {
  public:
    void Init(FEXCore::Context::ContextImpl* CTX);

    std::optional<CodeHitData> Lookup(Core::InternalThreadState* Thread, std::optional<ExecutableFileSectionInfo> Region, uint64_t GuestRIP,
                                      std::optional<uint64_t>& GuestCodeKey);
    void Validate(uint64_t GuestCodeKey, const CodeHitData& Hit, const CPU::CPUBackend::CompiledCode& CompiledCode,
                  std::optional<ExecutableFileSectionInfo> Region);
    bool Store(Core::InternalThreadState* Thread, std::optional<ExecutableFileSectionInfo> Region, uint64_t GuestRIP, uint64_t GuestCodeKey,
               std::span<const uint8_t> GuestCode, const CPU::CPUBackend::CompiledCode& CompiledCode,
               std::span<const FEXCore::CPU::Relocation> Relocations, const Frontend::Decoder::DecodedBlockInformation* DecodedBlockInfo);

    bool IsWritingDiskCache() const {
      return WritingDiskCache;
    }
    bool IsReadingDiskCache() const {
      return ReadingDiskCache;
    }
    bool IsValidating() const {
      return Validation;
    }

  private:
    bool OpenCacheDB(const fextl::string& CacheDBName, bool ReadOnly);
    uint64_t MakeLookupKey(Core::InternalThreadState* Thread, const uint64_t ModuleOffset, bool Writable, bool MonoBackpatcher);

    bool ReadingDiskCache {};
    bool WritingDiskCache {};
    FEXCore::Context::ContextImpl* CTX;
    XXH128_hash_t BucketHash;
    fextl::vector<fextl::unique_ptr<IndexedDB>> ROCacheDBs;
    fextl::unique_ptr<IndexedDB> RWCacheDB;
    Index Index;
    std::mutex IndexLock;
    bool FoundMetadata = false;
    struct CacheStoreWorkItem;

    // the Writer holds references to all this stuff above and needs to be last
    fextl::unique_ptr<WorkQueueThread> Writer;

    FEX_CONFIG_OPT(EnableDiskCache, DISKCACHE);
    FEX_CONFIG_OPT(Validation, DISKCACHEVALIDATION);
    FEX_CONFIG_OPT(MapDiskCacheFiles, DISKCACHEFILEMAPPING);
    FEX_CONFIG_OPT(RelocationFilter, DISKCACHERELOCATIONFILTER);
    FEX_CONFIG_OPT(AnonCaching, DISKCACHEANONCACHING);
    FEX_CONFIG_OPT(BasePathOverride, DISKCACHEPATH);
    FEX_CONFIG_OPT(RODBNames, DISKCACHERODBNAMES);
  };

  static constexpr uint16_t AnonPrefixGuestBytes = 64;
  // The current version of the diskcache.
  // This must be changed any time codegen changes occur!
  // Be aware of the impact of changing this frequently!
  static constexpr uint16_t FormatVersion = 24;

} // namespace DiskCache

} // namespace FEXCore