/**
 ******************************************************************************
 * Xenia : Xbox 360 Emulator Research Project                                 *
 ******************************************************************************
 * Copyright 2016 Ben Vanik. All rights reserved.                             *
 * Released under the BSD license - see LICENSE in the root for more details. *
 ******************************************************************************
 */

#include "xenia/gpu/vulkan/texture_cache.h"

#include "xenia/base/logging.h"
#include "xenia/base/math.h"
#include "xenia/base/memory.h"
#include "xenia/base/profiling.h"
#include "xenia/gpu/gpu_flags.h"
#include "xenia/gpu/sampler_info.h"
#include "xenia/gpu/texture_info.h"
#include "xenia/gpu/vulkan/vulkan_gpu_flags.h"
#include "xenia/ui/vulkan/vulkan_mem_alloc.h"

namespace xe {
namespace gpu {
namespace vulkan {

using xe::ui::vulkan::CheckResult;

constexpr uint32_t kMaxTextureSamplers = 32;
constexpr VkDeviceSize kStagingBufferSize = 32 * 1024 * 1024;

struct TextureConfig {
  VkFormat host_format;
};

static const TextureConfig texture_configs[64] = {
    /* k_1_REVERSE              */ {VK_FORMAT_UNDEFINED},
    /* k_1                      */ {VK_FORMAT_UNDEFINED},
    /* k_8                      */ {VK_FORMAT_R8_UNORM},
    // ! A1BGR5
    /* k_1_5_5_5                */ {VK_FORMAT_A1R5G5B5_UNORM_PACK16},
    /* k_5_6_5                  */ {VK_FORMAT_R5G6B5_UNORM_PACK16},
    /* k_6_5_5                  */ {VK_FORMAT_UNDEFINED},
    /* k_8_8_8_8                */ {VK_FORMAT_R8G8B8A8_UNORM},
    /* k_2_10_10_10             */ {VK_FORMAT_A2R10G10B10_UNORM_PACK32},
    /* k_8_A                    */ {VK_FORMAT_UNDEFINED},
    /* k_8_B                    */ {VK_FORMAT_UNDEFINED},
    /* k_8_8                    */ {VK_FORMAT_R8G8_UNORM},
    /* k_Cr_Y1_Cb_Y0            */ {VK_FORMAT_UNDEFINED},
    /* k_Y1_Cr_Y0_Cb            */ {VK_FORMAT_UNDEFINED},
    /* k_Shadow                 */ {VK_FORMAT_UNDEFINED},
    /* k_8_8_8_8_A              */ {VK_FORMAT_UNDEFINED},
    /* k_4_4_4_4                */ {VK_FORMAT_R4G4B4A4_UNORM_PACK16},
    // TODO: Verify if these two are correct (I think not).
    /* k_10_11_11               */ {VK_FORMAT_B10G11R11_UFLOAT_PACK32},
    /* k_11_11_10               */ {VK_FORMAT_B10G11R11_UFLOAT_PACK32},

    /* k_DXT1                   */ {VK_FORMAT_BC1_RGBA_UNORM_BLOCK},
    /* k_DXT2_3                 */ {VK_FORMAT_BC2_UNORM_BLOCK},
    /* k_DXT4_5                 */ {VK_FORMAT_BC3_UNORM_BLOCK},
    /* k_DXV                    */ {VK_FORMAT_UNDEFINED},

    // TODO: D24 unsupported on AMD.
    /* k_24_8                   */ {VK_FORMAT_D24_UNORM_S8_UINT},
    /* k_24_8_FLOAT             */ {VK_FORMAT_D32_SFLOAT_S8_UINT},
    /* k_16                     */ {VK_FORMAT_R16_UNORM},
    /* k_16_16                  */ {VK_FORMAT_R16G16_UNORM},
    /* k_16_16_16_16            */ {VK_FORMAT_R16G16B16A16_UNORM},
    /* k_16_EXPAND              */ {VK_FORMAT_R16_UNORM},
    /* k_16_16_EXPAND           */ {VK_FORMAT_R16G16_UNORM},
    /* k_16_16_16_16_EXPAND     */ {VK_FORMAT_R16G16B16A16_UNORM},
    /* k_16_FLOAT               */ {VK_FORMAT_R16_SFLOAT},
    /* k_16_16_FLOAT            */ {VK_FORMAT_R16G16_SFLOAT},
    /* k_16_16_16_16_FLOAT      */ {VK_FORMAT_R16G16B16A16_SFLOAT},

    // ! These are UNORM formats, not SINT.
    /* k_32                     */ {VK_FORMAT_R32_SINT},
    /* k_32_32                  */ {VK_FORMAT_R32G32_SINT},
    /* k_32_32_32_32            */ {VK_FORMAT_R32G32B32A32_SINT},
    /* k_32_FLOAT               */ {VK_FORMAT_R32_SFLOAT},
    /* k_32_32_FLOAT            */ {VK_FORMAT_R32G32_SFLOAT},
    /* k_32_32_32_32_FLOAT      */ {VK_FORMAT_R32G32B32A32_SFLOAT},
    /* k_32_AS_8                */ {VK_FORMAT_UNDEFINED},
    /* k_32_AS_8_8              */ {VK_FORMAT_UNDEFINED},
    /* k_16_MPEG                */ {VK_FORMAT_UNDEFINED},
    /* k_16_16_MPEG             */ {VK_FORMAT_UNDEFINED},
    /* k_8_INTERLACED           */ {VK_FORMAT_UNDEFINED},
    /* k_32_AS_8_INTERLACED     */ {VK_FORMAT_UNDEFINED},
    /* k_32_AS_8_8_INTERLACED   */ {VK_FORMAT_UNDEFINED},
    /* k_16_INTERLACED          */ {VK_FORMAT_UNDEFINED},
    /* k_16_MPEG_INTERLACED     */ {VK_FORMAT_UNDEFINED},
    /* k_16_16_MPEG_INTERLACED  */ {VK_FORMAT_UNDEFINED},

    // http://fileadmin.cs.lth.se/cs/Personal/Michael_Doggett/talks/unc-xenos-doggett.pdf
    /* k_DXN                    */ {VK_FORMAT_BC5_UNORM_BLOCK},  // ?

    /* k_8_8_8_8_AS_16_16_16_16 */ {VK_FORMAT_R8G8B8A8_UNORM},
    /* k_DXT1_AS_16_16_16_16    */ {VK_FORMAT_BC1_RGBA_UNORM_BLOCK},
    /* k_DXT2_3_AS_16_16_16_16  */ {VK_FORMAT_BC2_UNORM_BLOCK},
    /* k_DXT4_5_AS_16_16_16_16  */ {VK_FORMAT_BC3_UNORM_BLOCK},

    /* k_2_10_10_10_AS_16_16_16_16 */ {VK_FORMAT_A2R10G10B10_UNORM_PACK32},

    // TODO: Verify if these two are correct (I think not).
    /* k_10_11_11_AS_16_16_16_16 */ {VK_FORMAT_B10G11R11_UFLOAT_PACK32},  // ?
    /* k_11_11_10_AS_16_16_16_16 */ {VK_FORMAT_B10G11R11_UFLOAT_PACK32},  // ?
    /* k_32_32_32_FLOAT         */ {VK_FORMAT_R32G32B32_SFLOAT},
    /* k_DXT3A                  */ {VK_FORMAT_UNDEFINED},
    /* k_DXT5A                  */ {VK_FORMAT_UNDEFINED},  // ATI1N

    // http://fileadmin.cs.lth.se/cs/Personal/Michael_Doggett/talks/unc-xenos-doggett.pdf
    /* k_CTX1                   */ {VK_FORMAT_R8G8_UINT},

    /* k_DXT3A_AS_1_1_1_1       */ {VK_FORMAT_UNDEFINED},

    // Unused.
    /* kUnknown                 */ {VK_FORMAT_UNDEFINED},
    /* kUnknown                 */ {VK_FORMAT_UNDEFINED},
};

TextureCache::TextureCache(Memory* memory, RegisterFile* register_file,
                           TraceWriter* trace_writer,
                           ui::vulkan::VulkanDevice* device)
    : memory_(memory),
      register_file_(register_file),
      trace_writer_(trace_writer),
      device_(device),
      staging_buffer_(device, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                      kStagingBufferSize),
      wb_staging_buffer_(device, VK_BUFFER_USAGE_TRANSFER_DST_BIT,
                         kStagingBufferSize) {}

TextureCache::~TextureCache() { Shutdown(); }

VkResult TextureCache::Initialize() {
  VkResult status = VK_SUCCESS;

  // Descriptor pool used for all of our cached descriptors.
  VkDescriptorPoolSize pool_sizes[1];
  pool_sizes[0].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  pool_sizes[0].descriptorCount = 32768;
  descriptor_pool_ = std::make_unique<ui::vulkan::DescriptorPool>(
      *device_, 32768,
      std::vector<VkDescriptorPoolSize>(pool_sizes, std::end(pool_sizes)));

  wb_command_pool_ = std::make_unique<ui::vulkan::CommandBufferPool>(
      *device_, device_->queue_family_index());

  // Check some device limits
  // On low sampler counts: Rarely would we experience over 16 unique samplers.
  // This code could be refactored to scale up/down to the # of samplers.
  auto& limits = device_->device_info().properties.limits;
  if (limits.maxPerStageDescriptorSamplers < kMaxTextureSamplers ||
      limits.maxPerStageDescriptorSampledImages < kMaxTextureSamplers) {
    XELOGE(
        "Physical device is unable to support required number of sampled "
        "images! Expect instability! (maxPerStageDescriptorSamplers=%d, "
        "maxPerStageDescriptorSampledImages=%d)",
        limits.maxPerStageDescriptorSamplers,
        limits.maxPerStageDescriptorSampledImages);
    // assert_always();
  }

  // Create the descriptor set layout used for rendering.
  // We always have the same number of samplers but only some are used.
  // The shaders will alias the bindings to the 4 dimensional types.
  VkDescriptorSetLayoutBinding bindings[1];
  bindings[0].binding = 0;
  bindings[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  bindings[0].descriptorCount = kMaxTextureSamplers;
  bindings[0].stageFlags =
      VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT;
  bindings[0].pImmutableSamplers = nullptr;

  VkDescriptorSetLayoutCreateInfo descriptor_set_layout_info;
  descriptor_set_layout_info.sType =
      VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
  descriptor_set_layout_info.pNext = nullptr;
  descriptor_set_layout_info.flags = 0;
  descriptor_set_layout_info.bindingCount =
      static_cast<uint32_t>(xe::countof(bindings));
  descriptor_set_layout_info.pBindings = bindings;
  status =
      vkCreateDescriptorSetLayout(*device_, &descriptor_set_layout_info,
                                  nullptr, &texture_descriptor_set_layout_);
  if (status != VK_SUCCESS) {
    return status;
  }

  status = staging_buffer_.Initialize();
  if (status != VK_SUCCESS) {
    return status;
  }

  status = wb_staging_buffer_.Initialize();
  if (status != VK_SUCCESS) {
    return status;
  }

  // Create a memory allocator for textures.
  VmaVulkanFunctions vulkan_funcs = {};
  ui::vulkan::FillVMAVulkanFunctions(&vulkan_funcs);

  VmaAllocatorCreateInfo alloc_info = {
      0, *device_, *device_, 0, 0, nullptr, nullptr, 0, nullptr, &vulkan_funcs,
  };
  status = vmaCreateAllocator(&alloc_info, &mem_allocator_);
  if (status != VK_SUCCESS) {
    vkDestroyDescriptorSetLayout(*device_, texture_descriptor_set_layout_,
                                 nullptr);
    return status;
  }

  invalidated_textures_sets_[0].reserve(64);
  invalidated_textures_sets_[1].reserve(64);
  invalidated_textures_ = &invalidated_textures_sets_[0];

  device_queue_ = device_->AcquireQueue(device_->queue_family_index());
  return VK_SUCCESS;
}

void TextureCache::Shutdown() {
  if (device_queue_) {
    device_->ReleaseQueue(device_queue_, device_->queue_family_index());
  }

  // Free all textures allocated.
  ClearCache();
  Scavenge();

  if (mem_allocator_ != nullptr) {
    vmaDestroyAllocator(mem_allocator_);
    mem_allocator_ = nullptr;
  }
  vkDestroyDescriptorSetLayout(*device_, texture_descriptor_set_layout_,
                               nullptr);
}

TextureCache::TextureRegion* TextureCache::AllocateTextureRegion(
    Texture* texture, VkOffset3D region_offset, VkExtent3D region_size,
    VkFormatFeatureFlags required_flags) {
  // Create an image first.
  VkImageCreateInfo image_info = {};
  image_info.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
  switch (texture->texture_info.dimension) {
    case Dimension::k1D:
    case Dimension::k2D:
      image_info.imageType = VK_IMAGE_TYPE_2D;
      break;
    case Dimension::k3D:
      image_info.imageType = VK_IMAGE_TYPE_3D;
      break;
    case Dimension::kCube:
      image_info.imageType = VK_IMAGE_TYPE_2D;
      image_info.flags |= VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT;
      break;
    default:
      assert_unhandled_case(texture->texture_info.dimension);
      return nullptr;
  }

  assert_not_null(texture->texture_info.format_info());
  auto& config =
      texture_configs[int(texture->texture_info.format_info()->format)];
  VkFormat format = config.host_format;
  if (format == VK_FORMAT_UNDEFINED) {
    XELOGE(
        "Texture Cache: Attempted to allocate texture format %s, which is "
        "defined as VK_FORMAT_UNDEFINED!",
        texture->texture_info.format_info()->name);
    return nullptr;
  }

  image_info.tiling = VK_IMAGE_TILING_OPTIMAL;
  image_info.usage =
      VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT;

  // Check the device limits for the format before we create it.
  VkFormatProperties props;
  vkGetPhysicalDeviceFormatProperties(*device_, texture->format, &props);
  if ((props.optimalTilingFeatures & required_flags) != required_flags) {
    // Texture needs conversion on upload to a native format.
    XELOGE(
        "Texture Cache: Invalid usage flag specified on format %s (vk %d) "
        "(0x%.8X != 0x%.8X)",
        texture->texture_info.format_info()->name, texture->format,
        (props.optimalTilingFeatures & required_flags), required_flags);
  }

  if (texture->texture_info.dimension != Dimension::kCube &&
      props.optimalTilingFeatures & VK_FORMAT_FEATURE_COLOR_ATTACHMENT_BIT) {
    // Add color attachment usage if it's supported.
    image_info.usage |= VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
  } else if (props.optimalTilingFeatures &
             VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT) {
    // Add depth/stencil usage as well.
    image_info.usage |= VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT;
  }

  if (props.optimalTilingFeatures & VK_FORMAT_FEATURE_BLIT_DST_BIT) {
    image_info.usage |= VK_IMAGE_USAGE_TRANSFER_DST_BIT;
  }

  VkImageFormatProperties image_props;
  vkGetPhysicalDeviceImageFormatProperties(
      *device_, texture->format, image_info.imageType, image_info.tiling,
      image_info.usage, image_info.flags, &image_props);

  // TODO(DrChat): Actually check the image properties.

  image_info.format = texture->format;
  image_info.extent = region_size;
  image_info.mipLevels = texture->texture_info.mip_levels;
  image_info.arrayLayers = texture->texture_info.depth + 1;
  image_info.samples = VK_SAMPLE_COUNT_1_BIT;
  image_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
  image_info.queueFamilyIndexCount = 0;
  image_info.pQueueFamilyIndices = nullptr;
  image_info.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;

  VkImage image;

  VmaAllocation allocation;
  VmaAllocationCreateInfo vma_create_info = {
      0, VMA_MEMORY_USAGE_GPU_ONLY, 0, 0, 0, nullptr, nullptr,
  };
  VmaAllocationInfo vma_info = {};
  VkResult status =
      vmaCreateImage(mem_allocator_, &image_info, &vma_create_info, &image,
                     &allocation, &vma_info);
  if (status != VK_SUCCESS) {
    // Allocation failed.
    return nullptr;
  }

  TextureRegion* region = new TextureRegion();
  region->texture = texture;

  region->region_offset = region_offset;
  region->region_size = region_size;

  region->image = image;
  region->image_layout = image_info.initialLayout;
  region->usage_flags = image_info.usage;
  region->allocation = allocation;
  region->allocation_info = vma_info;

  region->region_contents_valid = false;
  return region;
}

TextureCache::Texture* TextureCache::AllocateTexture(
    const TextureInfo& texture_info, VkFormatFeatureFlags required_flags) {
  VkOffset3D region_offset = {0, 0, 0};
  VkExtent3D region_extent = {texture_info.width + 1, texture_info.height + 1,
                              1};

  auto texture = new Texture();
  texture->framebuffer = nullptr;
  texture->access_watch_handle = 0;
  texture->texture_info = texture_info;

  assert_not_null(texture->texture_info.format_info());
  auto& config =
      texture_configs[int(texture->texture_info.format_info()->format)];
  texture->format = config.host_format != VK_FORMAT_UNDEFINED
                        ? config.host_format
                        : VK_FORMAT_R8G8B8A8_UNORM;

  TextureRegion* base_region = AllocateTextureRegion(
      texture, region_offset, region_extent, required_flags);
  if (!base_region) {
    delete texture;
    return nullptr;
  }

  texture->base_region = base_region;
  return texture;
}

bool TextureCache::FreeTexture(Texture* texture) {
  if (texture->in_flight_fence) {
    VkResult status = vkGetFenceStatus(*device_, texture->in_flight_fence);
    if (status != VK_SUCCESS && status != VK_ERROR_DEVICE_LOST) {
      // Texture still in flight.
      return false;
    }
  }

  if (texture->framebuffer) {
    vkDestroyFramebuffer(*device_, texture->framebuffer, nullptr);
  }

  for (auto region_it = texture->regions.begin();
       region_it != texture->regions.end(); ++region_it) {
    TextureRegion* region = region_it->get();
    for (auto& view : region->views) {
      vkDestroyImageView(*device_, view->view, nullptr);
    }
    vmaDestroyImage(mem_allocator_, region->image, region->allocation);
  }
  texture->regions.clear();

  // Free the base region (which is not part of regions)
  for (auto& view : texture->base_region->views) {
    vkDestroyImageView(*device_, view->view, nullptr);
  }
  vmaDestroyImage(mem_allocator_, texture->base_region->image,
                  texture->base_region->allocation);

  if (texture->access_watch_handle) {
    memory_->CancelAccessWatch(texture->access_watch_handle);
    texture->access_watch_handle = 0;
  }

  delete texture;
  return true;
}

void TextureCache::WatchCallback(void* context_ptr, void* data_ptr,
                                 uint32_t address) {
  auto self = reinterpret_cast<TextureCache*>(context_ptr);
  auto touched_texture = reinterpret_cast<Texture*>(data_ptr);
  // Clear watch handle first so we don't redundantly
  // remove.
  assert_not_zero(touched_texture->access_watch_handle);
  touched_texture->access_watch_handle = 0;
  touched_texture->pending_invalidation = true;

  // Add to pending list so Scavenge will clean it up.
  self->invalidated_textures_mutex_.lock();
  self->invalidated_textures_->insert(touched_texture);
  self->invalidated_textures_mutex_.unlock();
}

TextureCache::Texture* TextureCache::DemandResolveTexture(
    const TextureInfo& texture_info) {
  auto texture_hash = texture_info.hash();
  for (auto it = textures_.find(texture_hash); it != textures_.end(); ++it) {
    if (it->second->texture_info == texture_info) {
      if (it->second->pending_invalidation) {
        // This texture has been invalidated!
        RemoveInvalidatedTextures();
        break;
      }

      for (auto& region : it->second->regions) {
        // Invalidate region contents
        region->region_contents_valid = false;
      }

      // Tell the trace writer to "cache" this memory (but not read it)
      trace_writer_->WriteMemoryReadCachedNop(texture_info.guest_address,
                                              texture_info.input_length);

      return it->second;
    }
  }

  VkFormatFeatureFlags required_flags = VK_FORMAT_FEATURE_SAMPLED_IMAGE_BIT;
  if (texture_info.texture_format == TextureFormat::k_24_8 ||
      texture_info.texture_format == TextureFormat::k_24_8_FLOAT) {
    required_flags |= VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT;
  } else {
    required_flags |= VK_FORMAT_FEATURE_COLOR_ATTACHMENT_BIT;
  }

  // No texture at this location. Make a new one.
  auto texture = AllocateTexture(texture_info, required_flags);
  if (!texture) {
    // Failed to allocate texture (out of memory)
    XELOGE("Vulkan Texture Cache: Failed to allocate texture!");
    return nullptr;
  }

  // Setup a debug name for the texture.
  device_->DbgSetObjectName(
      reinterpret_cast<uint64_t>(texture->base_region->image),
      VK_DEBUG_REPORT_OBJECT_TYPE_IMAGE_EXT,
      xe::format_string(
          "0x%.8X - 0x%.8X", texture_info.guest_address,
          texture_info.guest_address + texture_info.input_length));

  // Setup an access watch. If this texture is touched, it is destroyed.
  texture->access_watch_handle = memory_->AddPhysicalAccessWatch(
      texture_info.guest_address, texture_info.input_length,
      cpu::MMIOHandler::kWatchWrite, &WatchCallback, this, texture);

  textures_[texture_hash] = texture;
  COUNT_profile_set("gpu/texture_cache/textures", textures_.size());
  return texture;
}

TextureCache::TextureRegion* TextureCache::DemandRegion(
    const TextureInfo& texture_info, VkCommandBuffer command_buffer,
    VkCommandBuffer setup_buffer, VkFence completion_fence) {
  // Run a tight loop to scan for an exact match existing texture.
  auto texture_hash = texture_info.hash();
  for (auto it = textures_.find(texture_hash); it != textures_.end(); ++it) {
    if (it->second->texture_info == texture_info) {
      if (it->second->pending_invalidation) {
        // This texture has been invalidated!
        RemoveInvalidatedTextures();
        break;
      }

      trace_writer_->WriteMemoryReadCached(texture_info.guest_address,
                                           texture_info.input_length);

      return it->second->base_region;
    }
  }

  // If we didn't find an exact match, see if we can find a subregion of an
  // existing texture to use.
  VkOffset2D offset;
  auto containing_tex = LookupAddress(
      texture_info.guest_address, texture_info.width + 1,
      texture_info.height + 1, texture_info.format_info()->format, &offset);
  if (containing_tex != nullptr) {
    // Found a larger texture that contains the requested texels.
    // Find/create/update a matching subregion.

    // First, see if the region exists
    TextureRegion* region = nullptr;
    for (auto region_it = containing_tex->regions.begin();
         region_it != containing_tex->regions.end(); ++region_it) {
      const VkOffset3D& roffset = (*region_it)->region_offset;
      const VkExtent3D& rsize = (*region_it)->region_size;
      if (roffset.x == offset.x && roffset.y == offset.y &&
          rsize.width == (texture_info.width + 1) &&
          rsize.height == (texture_info.height + 1)) {
        // Match found
        region = region_it->get();
        break;
      }
    }

    if (!region) {
      // No match, create a new region and add it the texture
      region = AllocateTextureRegion(
          containing_tex, {offset.x, offset.y, 0},
          {texture_info.width + 1, texture_info.height + 1, 1},
          VK_FORMAT_FEATURE_SAMPLED_IMAGE_BIT);
      containing_tex->regions.push_back(std::unique_ptr<TextureRegion>(region));
    }

    if (command_buffer && region && !region->region_contents_valid) {
      // Region content is out-of-date, recreate it by blitting from the base
      // region.

      // Transition the region into a transfer destination layout.
      VkImageMemoryBarrier barrier;
      barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
      barrier.pNext = nullptr;
      barrier.srcAccessMask = 0;
      barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
      barrier.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
      barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
      barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
      barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
      barrier.image = region->image;
      barrier.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};
      if (containing_tex->format == VK_FORMAT_D16_UNORM_S8_UINT ||
          containing_tex->format == VK_FORMAT_D24_UNORM_S8_UINT ||
          containing_tex->format == VK_FORMAT_D32_SFLOAT_S8_UINT) {
        barrier.subresourceRange.aspectMask =
            VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT;
      }

      vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
                           VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, 0, 0, nullptr, 0,
                           nullptr, 1, &barrier);

      // Now move the converted texture into the destination.
      VkImageCopy copy_region;

      copy_region.srcSubresource = {barrier.subresourceRange.aspectMask, 0, 0,
                                    1};
      copy_region.dstSubresource = {barrier.subresourceRange.aspectMask, 0, 0,
                                    1};
      copy_region.srcOffset = {offset.x, offset.y, 0};
      copy_region.dstOffset = {0, 0, 0};
      copy_region.extent = {texture_info.width + 1, texture_info.height + 1, 1};

      vkCmdCopyImage(command_buffer, containing_tex->base_region->image,
                     containing_tex->base_region->image_layout, region->image,
                     VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &copy_region);

      // Now transition the region into a shader readonly source.
      barrier.srcAccessMask = barrier.dstAccessMask;
      barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
      barrier.oldLayout = barrier.newLayout;
      barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
      vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
                           VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, 0, 0, nullptr, 0,
                           nullptr, 1, &barrier);

      region->image_layout = barrier.newLayout;

      // Mark contents as up-to-date
      region->region_contents_valid = true;
    }

    return region;
  }

  // Create a new texture and cache it.
  auto texture =
      AllocateTexture(texture_info, VK_FORMAT_FEATURE_SAMPLED_IMAGE_BIT);
  if (!texture) {
    // Failed to allocate texture (out of memory)
    XELOGE("Vulkan Texture Cache: Failed to allocate texture!");
    return nullptr;
  }

  trace_writer_->WriteMemoryRead(texture_info.guest_address,
                                 texture_info.input_length);

  // Okay. Put a writewatch on it to tell us if it's been modified from the
  // guest.
  texture->access_watch_handle = memory_->AddPhysicalAccessWatch(
      texture_info.guest_address, texture_info.input_length,
      cpu::MMIOHandler::kWatchWrite, &WatchCallback, this, texture);

  if (!UploadTexture(setup_buffer, completion_fence, texture, texture_info)) {
    FreeTexture(texture);
    return nullptr;
  }

  // Setup a debug name for the texture.
  device_->DbgSetObjectName(
      reinterpret_cast<uint64_t>(texture->base_region->image),
      VK_DEBUG_REPORT_OBJECT_TYPE_IMAGE_EXT,
      xe::format_string(
          "0x%.8X - 0x%.8X", texture_info.guest_address,
          texture_info.guest_address + texture_info.input_length));

  textures_[texture_hash] = texture;
  COUNT_profile_set("gpu/texture_cache/textures", textures_.size());
  return texture->base_region;
}

TextureCache::TextureRegionView* TextureCache::DemandTextureRegionView(
    TextureRegion* region, uint16_t swizzle) {
  for (auto it = region->views.begin(); it != region->views.end(); ++it) {
    if ((*it)->swizzle == swizzle) {
      return (*it).get();
    }
  }

  VkImageViewCreateInfo view_info;
  view_info.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
  view_info.pNext = nullptr;
  view_info.flags = 0;
  view_info.image = region->image;
  view_info.format = region->texture->format;

  switch (region->texture->texture_info.dimension) {
    case Dimension::k1D:
    case Dimension::k2D:
      view_info.viewType = VK_IMAGE_VIEW_TYPE_2D;
      break;
    case Dimension::k3D:
      view_info.viewType = VK_IMAGE_VIEW_TYPE_3D;
      break;
    case Dimension::kCube:
      view_info.viewType = VK_IMAGE_VIEW_TYPE_CUBE;
      break;
    default:
      assert_always();
  }

  VkComponentSwizzle swiz_component_map[] = {
      VK_COMPONENT_SWIZZLE_R,        VK_COMPONENT_SWIZZLE_G,
      VK_COMPONENT_SWIZZLE_B,        VK_COMPONENT_SWIZZLE_A,
      VK_COMPONENT_SWIZZLE_ZERO,     VK_COMPONENT_SWIZZLE_ONE,
      VK_COMPONENT_SWIZZLE_IDENTITY,
  };
  if (region->texture->texture_info.texture_format ==
      TextureFormat::k_4_4_4_4) {
    swiz_component_map[0] = VK_COMPONENT_SWIZZLE_A;
    swiz_component_map[1] = VK_COMPONENT_SWIZZLE_B;
    swiz_component_map[2] = VK_COMPONENT_SWIZZLE_G;
    swiz_component_map[3] = VK_COMPONENT_SWIZZLE_R;
  }

  view_info.components = {
      swiz_component_map[(swizzle >> 0) & 0x7],
      swiz_component_map[(swizzle >> 3) & 0x7],
      swiz_component_map[(swizzle >> 6) & 0x7],
      swiz_component_map[(swizzle >> 9) & 0x7],
  };
  view_info.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0,
                                region->texture->texture_info.mip_levels, 0, 1};
  if (region->texture->format == VK_FORMAT_D16_UNORM_S8_UINT ||
      region->texture->format == VK_FORMAT_D24_UNORM_S8_UINT ||
      region->texture->format == VK_FORMAT_D32_SFLOAT_S8_UINT) {
    // This applies to any depth/stencil format, but we only use D24S8 / D32FS8.
    view_info.subresourceRange.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;
  }

  if (region->texture->texture_info.dimension == Dimension::kCube) {
    view_info.subresourceRange.layerCount = 6;
  }

  VkImageView view;
  auto status = vkCreateImageView(*device_, &view_info, nullptr, &view);
  CheckResult(status, "vkCreateImageView");
  if (status == VK_SUCCESS) {
    auto region_view = new TextureRegionView();
    region_view->region = region;
    region_view->view = view;
    region_view->swizzle = swizzle;
    region->views.push_back(std::unique_ptr<TextureRegionView>(region_view));
    return region_view;
  }

  return nullptr;
}

TextureCache::Sampler* TextureCache::Demand(const SamplerInfo& sampler_info) {
#if FINE_GRAINED_DRAW_SCOPES
  SCOPE_profile_cpu_f("gpu");
#endif  // FINE_GRAINED_DRAW_SCOPES

  auto sampler_hash = sampler_info.hash();
  for (auto it = samplers_.find(sampler_hash); it != samplers_.end(); ++it) {
    if (it->second->sampler_info == sampler_info) {
      // Found a compatible sampler.
      return it->second;
    }
  }

  VkResult status = VK_SUCCESS;

  // Create a new sampler and cache it.
  VkSamplerCreateInfo sampler_create_info;
  sampler_create_info.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
  sampler_create_info.pNext = nullptr;
  sampler_create_info.flags = 0;
  sampler_create_info.mipmapMode = VK_SAMPLER_MIPMAP_MODE_NEAREST;
  sampler_create_info.maxAnisotropy = 1.0f;

  // Texture level filtering.
  VkSamplerMipmapMode mip_filter;
  switch (sampler_info.mip_filter) {
    case TextureFilter::kBaseMap:
      // TODO(DrChat): ?
      mip_filter = VK_SAMPLER_MIPMAP_MODE_NEAREST;
      break;
    case TextureFilter::kPoint:
      mip_filter = VK_SAMPLER_MIPMAP_MODE_NEAREST;
      break;
    case TextureFilter::kLinear:
      mip_filter = VK_SAMPLER_MIPMAP_MODE_LINEAR;
      break;
    default:
      assert_unhandled_case(sampler_info.mip_filter);
      return nullptr;
  }

  VkFilter min_filter;
  switch (sampler_info.min_filter) {
    case TextureFilter::kPoint:
      min_filter = VK_FILTER_NEAREST;
      break;
    case TextureFilter::kLinear:
      min_filter = VK_FILTER_LINEAR;
      break;
    default:
      assert_unhandled_case(sampler_info.min_filter);
      return nullptr;
  }
  VkFilter mag_filter;
  switch (sampler_info.mag_filter) {
    case TextureFilter::kPoint:
      mag_filter = VK_FILTER_NEAREST;
      break;
    case TextureFilter::kLinear:
      mag_filter = VK_FILTER_LINEAR;
      break;
    default:
      assert_unhandled_case(mag_filter);
      return nullptr;
  }

  sampler_create_info.minFilter = min_filter;
  sampler_create_info.magFilter = mag_filter;
  sampler_create_info.mipmapMode = mip_filter;

  // FIXME: Both halfway / mirror clamp to border aren't mapped properly.
  VkSamplerAddressMode address_mode_map[] = {
      /* kRepeat               */ VK_SAMPLER_ADDRESS_MODE_REPEAT,
      /* kMirroredRepeat       */ VK_SAMPLER_ADDRESS_MODE_MIRRORED_REPEAT,
      /* kClampToEdge          */ VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
      /* kMirrorClampToEdge    */ VK_SAMPLER_ADDRESS_MODE_MIRROR_CLAMP_TO_EDGE,
      /* kClampToHalfway       */ VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
      /* kMirrorClampToHalfway */ VK_SAMPLER_ADDRESS_MODE_MIRROR_CLAMP_TO_EDGE,
      /* kClampToBorder        */ VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER,
      /* kMirrorClampToBorder  */ VK_SAMPLER_ADDRESS_MODE_MIRROR_CLAMP_TO_EDGE,
  };
  sampler_create_info.addressModeU =
      address_mode_map[static_cast<int>(sampler_info.clamp_u)];
  sampler_create_info.addressModeV =
      address_mode_map[static_cast<int>(sampler_info.clamp_v)];
  sampler_create_info.addressModeW =
      address_mode_map[static_cast<int>(sampler_info.clamp_w)];

  float aniso = 0.f;
  switch (sampler_info.aniso_filter) {
    case AnisoFilter::kDisabled:
      aniso = 1.0f;
      break;
    case AnisoFilter::kMax_1_1:
      aniso = 1.0f;
      break;
    case AnisoFilter::kMax_2_1:
      aniso = 2.0f;
      break;
    case AnisoFilter::kMax_4_1:
      aniso = 4.0f;
      break;
    case AnisoFilter::kMax_8_1:
      aniso = 8.0f;
      break;
    case AnisoFilter::kMax_16_1:
      aniso = 16.0f;
      break;
    default:
      assert_unhandled_case(aniso);
      return nullptr;
  }

  sampler_create_info.anisotropyEnable =
      sampler_info.aniso_filter != AnisoFilter::kDisabled ? VK_TRUE : VK_FALSE;
  sampler_create_info.maxAnisotropy = aniso;

  sampler_create_info.compareEnable = VK_FALSE;
  sampler_create_info.compareOp = VK_COMPARE_OP_NEVER;
  sampler_create_info.mipLodBias = sampler_info.lod_bias;
  sampler_create_info.minLod = float(sampler_info.mip_min_level);
  sampler_create_info.maxLod = float(sampler_info.mip_max_level);
  sampler_create_info.borderColor = VK_BORDER_COLOR_FLOAT_TRANSPARENT_BLACK;
  sampler_create_info.unnormalizedCoordinates = VK_FALSE;
  VkSampler vk_sampler;
  status =
      vkCreateSampler(*device_, &sampler_create_info, nullptr, &vk_sampler);
  CheckResult(status, "vkCreateSampler");
  if (status != VK_SUCCESS) {
    return nullptr;
  }

  auto sampler = new Sampler();
  sampler->sampler = vk_sampler;
  sampler->sampler_info = sampler_info;
  samplers_[sampler_hash] = sampler;

  return sampler;
}

bool TextureFormatIsSimilar(TextureFormat left, TextureFormat right) {
#define COMPARE_FORMAT(x, y)                                     \
  if ((left == TextureFormat::x && right == TextureFormat::y) || \
      (left == TextureFormat::y && right == TextureFormat::x)) { \
    return true;                                                 \
  }

  if (left == right) return true;
  if (GetBaseFormat(left) == GetBaseFormat(right)) return true;

  return false;
#undef COMPARE_FORMAT
}

TextureCache::Texture* TextureCache::Lookup(const TextureInfo& texture_info) {
  auto texture_hash = texture_info.hash();
  for (auto it = textures_.find(texture_hash); it != textures_.end(); ++it) {
    if (it->second->texture_info == texture_info) {
      return it->second;
    }
  }
  // slow path
  for (auto it = textures_.begin(); it != textures_.end(); ++it) {
    const auto& other_texture_info = it->second->texture_info;
#define COMPARE_FIELD(x) \
  if (texture_info.x != other_texture_info.x) continue
    COMPARE_FIELD(guest_address);
    COMPARE_FIELD(dimension);
    COMPARE_FIELD(width);
    COMPARE_FIELD(height);
    COMPARE_FIELD(depth);
    COMPARE_FIELD(endianness);
    COMPARE_FIELD(is_tiled);
    COMPARE_FIELD(has_packed_mips);
    COMPARE_FIELD(input_length);
#undef COMPARE_FIELD
    if (!TextureFormatIsSimilar(texture_info.texture_format,
                                other_texture_info.texture_format)) {
      continue;
    }
    /*const auto format_info = texture_info.format_info();
    const auto other_format_info = other_texture_info.format_info();
#define COMPARE_FIELD(x) if (format_info->x != other_format_info->x) continue
    COMPARE_FIELD(type);
    COMPARE_FIELD(block_width);
    COMPARE_FIELD(block_height);
    COMPARE_FIELD(bits_per_pixel);
#undef COMPARE_FIELD*/
    return it->second;
  }

  return nullptr;
}

TextureCache::Texture* TextureCache::LookupAddress(uint32_t guest_address,
                                                   uint32_t width,
                                                   uint32_t height,
                                                   TextureFormat format,
                                                   VkOffset2D* out_offset) {
  for (auto it = textures_.begin(); it != textures_.end(); ++it) {
    const auto& texture_info = it->second->texture_info;
    if (guest_address >= texture_info.guest_address &&
        guest_address <
            texture_info.guest_address + texture_info.input_length &&
        texture_info.size_2d.input_width >= width &&
        texture_info.size_2d.input_height >= height && out_offset) {
      auto offset_bytes = guest_address - texture_info.guest_address;

      if (texture_info.dimension == Dimension::k2D) {
        out_offset->x = 0;
        out_offset->y = offset_bytes / texture_info.size_2d.input_pitch;
        if (offset_bytes % texture_info.size_2d.input_pitch != 0) {
          // TODO: offset_x
        }
      }

      return it->second;
    }

    if (texture_info.guest_address == guest_address &&
        texture_info.dimension == Dimension::k2D &&
        texture_info.size_2d.input_width == width &&
        texture_info.size_2d.input_height == height) {
      if (out_offset) {
        out_offset->x = 0;
        out_offset->y = 0;
      }

      return it->second;
    }
  }

  return nullptr;
}

void TextureSwap(Endian endianness, void* dest, const void* src,
                 size_t length) {
  switch (endianness) {
    case Endian::k8in16:
      xe::copy_and_swap_16_unaligned(dest, src, length / 2);
      break;
    case Endian::k8in32:
      xe::copy_and_swap_32_unaligned(dest, src, length / 4);
      break;
    case Endian::k16in32:  // Swap high and low 16 bits within a 32 bit word
      xe::copy_and_swap_16_in_32_unaligned(dest, src, length);
      break;
    default:
    case Endian::kUnspecified:
      std::memcpy(dest, src, length);
      break;
  }
}

void TextureCache::FlushPendingCommands(VkCommandBuffer setup_buffer,
                                        VkFence completion_fence) {
  auto status = vkEndCommandBuffer(setup_buffer);
  CheckResult(status, "vkEndCommandBuffer");

  VkSubmitInfo submit_info;
  std::memset(&submit_info, 0, sizeof(submit_info));
  submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
  submit_info.commandBufferCount = 1;
  submit_info.pCommandBuffers = &setup_buffer;

  if (device_queue_) {
    auto status =
        vkQueueSubmit(device_queue_, 1, &submit_info, completion_fence);
    CheckResult(status, "vkQueueSubmit");
  } else {
    std::lock_guard<std::mutex>(device_->primary_queue_mutex());

    auto status = vkQueueSubmit(device_->primary_queue(), 1, &submit_info,
                                completion_fence);
    CheckResult(status, "vkQueueSubmit");
  }

  vkWaitForFences(*device_, 1, &completion_fence, VK_TRUE, -1);
  staging_buffer_.Scavenge();
  vkResetFences(*device_, 1, &completion_fence);

  // Reset the command buffer and put it back into the recording state.
  vkResetCommandBuffer(setup_buffer, 0);
  VkCommandBufferBeginInfo begin_info;
  std::memset(&begin_info, 0, sizeof(begin_info));
  begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
  begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
  vkBeginCommandBuffer(setup_buffer, &begin_info);
}

bool TextureCache::ConvertTexture2D(uint8_t* dest,
                                    VkBufferImageCopy* copy_region,
                                    uint32_t mip, const TextureInfo& src) {
  uint32_t offset_x = 0;
  uint32_t offset_y = 0;
  uint32_t address =
      TextureInfo::GetMipLocation(src, mip, &offset_x, &offset_y);
  void* host_address = memory_->TranslatePhysical(address);

  // Pitch of the source texture in blocks.
  uint32_t block_width = mip == 0
                             ? src.size_2d.block_width
                             : xe::next_pow2(src.size_2d.block_width) >> mip;
  uint32_t logical_width = src.size_2d.logical_width >> mip;
  uint32_t logical_height = src.size_2d.logical_height >> mip;
  uint32_t input_width = src.size_2d.input_width >> mip;
  uint32_t input_height = src.size_2d.input_height >> mip;

  // All dimensions must be a multiple of block w/h
  logical_width = xe::round_up(logical_width, src.format_info()->block_width);
  logical_height =
      xe::round_up(logical_height, src.format_info()->block_height);
  input_width = xe::round_up(input_width, src.format_info()->block_width);
  input_height = xe::round_up(input_height, src.format_info()->block_height);

  if (!src.is_tiled) {
    uint32_t bytes_per_block = src.format_info()->block_width *
                               src.format_info()->block_height *
                               src.format_info()->bits_per_pixel / 8;

    const uint8_t* src_mem = reinterpret_cast<const uint8_t*>(host_address);
    src_mem += offset_y * src.size_2d.input_pitch;
    src_mem += offset_x * bytes_per_block;
    TextureSwap(src.endianness, dest, src_mem,
                src.size_2d.input_pitch * src.size_2d.logical_height);
  } else {
    // Untile image.
    // We could do this in a shader to speed things up, as this is pretty
    // slow.
    const uint8_t* src_mem = reinterpret_cast<const uint8_t*>(host_address);
    TextureInfo::ConvertTiled(dest, src_mem, src.endianness, src.format_info(),
                              offset_x, offset_y, block_width, logical_width,
                              logical_height, input_width);
  }

  copy_region->bufferRowLength = input_width;
  copy_region->bufferImageHeight = input_height;
  copy_region->imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, mip, 0, 1};
  copy_region->imageExtent = {logical_width, logical_height, 1};
  return true;
}

bool TextureCache::ConvertTextureCube(uint8_t* dest,
                                      VkBufferImageCopy* copy_region,
                                      const TextureInfo& src) {
  void* host_address = memory_->TranslatePhysical(src.guest_address);
  if (!src.is_tiled) {
    // Fast path copy entire image.
    TextureSwap(src.endianness, dest, host_address, src.input_length);
    copy_region->bufferRowLength = src.size_cube.input_width;
    copy_region->bufferImageHeight = src.size_cube.input_height;
    copy_region->imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 6};
    copy_region->imageExtent = {src.size_cube.logical_width,
                                src.size_cube.logical_height, 1};
    return true;
  } else {
    // TODO(benvanik): optimize this inner loop (or work by tiles).
    const uint8_t* src_mem = reinterpret_cast<const uint8_t*>(host_address);
    uint32_t bytes_per_block = src.format_info()->block_width *
                               src.format_info()->block_height *
                               src.format_info()->bits_per_pixel / 8;
    // Tiled textures can be packed; get the offset into the packed texture.
    uint32_t offset_x;
    uint32_t offset_y;
    TextureInfo::GetPackedTileOffset(src, &offset_x, &offset_y);
    auto bpp = (bytes_per_block >> 2) +
               ((bytes_per_block >> 1) >> (bytes_per_block >> 2));
    for (int face = 0; face < 6; ++face) {
      for (uint32_t y = 0, output_base_offset = 0;
           y < src.size_cube.block_height;
           y++, output_base_offset += src.size_cube.input_pitch) {
        auto input_base_offset = TextureInfo::TiledOffset2DOuter(
            offset_y + y,
            (src.size_cube.input_width / src.format_info()->block_width), bpp);
        for (uint32_t x = 0, output_offset = output_base_offset;
             x < src.size_cube.block_width;
             x++, output_offset += bytes_per_block) {
          auto input_offset =
              TextureInfo::TiledOffset2DInner(offset_x + x, offset_y + y, bpp,
                                              input_base_offset) >>
              bpp;
          TextureSwap(src.endianness, dest + output_offset,
                      src_mem + input_offset * bytes_per_block,
                      bytes_per_block);
        }
      }
      src_mem += src.size_cube.input_face_length;
      dest += src.size_cube.input_face_length;
    }

    copy_region->bufferRowLength = src.size_cube.input_width;
    copy_region->bufferImageHeight = src.size_cube.input_height;
    copy_region->imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 6};
    copy_region->imageExtent = {src.size_cube.logical_width,
                                src.size_cube.logical_height, 1};
    return true;
  }

  return false;
}

bool TextureCache::ConvertTexture(uint8_t* dest, VkBufferImageCopy* copy_region,
                                  uint32_t mip, const TextureInfo& src) {
  switch (src.dimension) {
    case Dimension::k1D:
      assert_always();
      break;
    case Dimension::k2D:
      return ConvertTexture2D(dest, copy_region, mip, src);
    case Dimension::k3D:
      assert_always();
      break;
    case Dimension::kCube:
      return ConvertTextureCube(dest, copy_region, src);
  }
  return false;
}

bool TextureCache::UploadTexture(VkCommandBuffer command_buffer,
                                 VkFence completion_fence, Texture* dest,
                                 const TextureInfo& src) {
#if FINE_GRAINED_DRAW_SCOPES
  SCOPE_profile_cpu_f("gpu");
#endif  // FINE_GRAINED_DRAW_SCOPES

  size_t unpack_length;
  if (!ComputeTextureStorage(&unpack_length, src)) {
    XELOGW("Failed to compute texture storage");
    return false;
  }

  size_t total_unpack_length = unpack_length;
  for (uint32_t i = 1; i < src.mip_levels; i++) {
    // Add in more space for mips.
    total_unpack_length += TextureInfo::GetMipLinearSize(src, i);
  }

  if (!staging_buffer_.CanAcquire(total_unpack_length)) {
    // Need to have unique memory for every upload for at least one frame. If we
    // run out of memory, we need to flush all queued upload commands to the
    // GPU.
    FlushPendingCommands(command_buffer, completion_fence);

    // Uploads have been flushed. Continue.
    if (!staging_buffer_.CanAcquire(total_unpack_length)) {
      // The staging buffer isn't big enough to hold this texture.
      XELOGE(
          "TextureCache staging buffer is too small! (uploading 0x%.8X bytes)",
          total_unpack_length);
      assert_always();
      return false;
    }
  }

  // Grab some temporary memory for staging.
  auto alloc = staging_buffer_.Acquire(total_unpack_length, completion_fence);
  assert_not_null(alloc);
  if (!alloc) {
    XELOGE("%s: Failed to acquire staging memory", __func__);
    return false;
  }

  // DEBUG: Check the source address. If it's completely zero'd out, print it.
  bool valid = false;
  auto src_data = memory_->TranslatePhysical(src.guest_address);
  for (uint32_t i = 0; i < src.input_length; i++) {
    if (src_data[i] != 0) {
      valid = true;
      break;
    }
  }

  if (!valid) {
    XELOGW(
        "Warning: Uploading blank texture at address 0x%.8X "
        "(length: 0x%.8X, format: %s)",
        src.guest_address, src.input_length, src.format_info()->name);
  }

  // Invalidate contents of all regions for this texture, except for the base
  // region.
  for (auto region_it = dest->regions.begin(); region_it != dest->regions.end();
       ++region_it) {
    (*region_it)->region_contents_valid = false;
  }
  dest->base_region->region_contents_valid = true;

  // Upload texture into GPU memory.
  // TODO: If the GPU supports it, we can submit a compute batch to convert the
  // texture and copy it to its destination. Otherwise, fallback to conversion
  // on the CPU.
  std::vector<VkBufferImageCopy> copy_regions(src.mip_levels);

  // Base MIP
  if (!ConvertTexture(reinterpret_cast<uint8_t*>(alloc->host_ptr),
                      &copy_regions[0], 0, src)) {
    XELOGW("Failed to convert texture");
    return false;
  }
  copy_regions[0].bufferOffset = alloc->offset;
  copy_regions[0].imageOffset = {0, 0, 0};

  // Now upload all the MIPs
  VkDeviceSize buffer_offset = unpack_length;
  for (uint32_t mip = 1; mip < src.mip_levels; mip++) {
    uint8_t* dest = reinterpret_cast<uint8_t*>(alloc->host_ptr) + buffer_offset;
    ConvertTexture(dest, &copy_regions[mip], mip, src);
    copy_regions[mip].bufferOffset = alloc->offset + buffer_offset;
    copy_regions[mip].imageOffset = {0, 0, 0};

    // With each mip, the length is divided by 4.
    buffer_offset += TextureInfo::GetMipLinearSize(src, mip);
  }

  // Transition the texture into a transfer destination layout.
  VkImageMemoryBarrier barrier;
  barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
  barrier.pNext = nullptr;
  barrier.srcAccessMask = 0;
  barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
  barrier.oldLayout = dest->base_region->image_layout;
  barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
  barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.image = dest->base_region->image;
  barrier.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, src.mip_levels,
                              copy_regions[0].imageSubresource.baseArrayLayer,
                              copy_regions[0].imageSubresource.layerCount};
  if (dest->format == VK_FORMAT_D16_UNORM_S8_UINT ||
      dest->format == VK_FORMAT_D24_UNORM_S8_UINT ||
      dest->format == VK_FORMAT_D32_SFLOAT_S8_UINT) {
    barrier.subresourceRange.aspectMask =
        VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT;
  }

  vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
                       VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0,
                       nullptr, 1, &barrier);

  // Now move the converted texture into the destination.
  if (dest->format == VK_FORMAT_D16_UNORM_S8_UINT ||
      dest->format == VK_FORMAT_D24_UNORM_S8_UINT ||
      dest->format == VK_FORMAT_D32_SFLOAT_S8_UINT) {
    // Do just a depth upload (for now).
    // This assumes depth buffers don't have mips (hopefully they don't)
    copy_regions[0].imageSubresource.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;
  }
  vkCmdCopyBufferToImage(command_buffer, staging_buffer_.gpu_buffer(),
                         dest->base_region->image,
                         VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, src.mip_levels,
                         copy_regions.data());

  // Now transition the texture into a shader readonly source.
  barrier.srcAccessMask = barrier.dstAccessMask;
  barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
  barrier.oldLayout = barrier.newLayout;
  barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
  vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_TRANSFER_BIT,
                       VK_PIPELINE_STAGE_VERTEX_SHADER_BIT |
                           VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
                       0, 0, nullptr, 0, nullptr, 1, &barrier);

  dest->base_region->image_layout = barrier.newLayout;
  return true;
}

bool TextureCache::ComputeTextureStorage(size_t* output_length,
                                         const TextureInfo& src) {
  if (src.texture_format == TextureFormat::k_CTX1) {
    switch (src.dimension) {
      case Dimension::k1D: {
        assert_always();
      } break;
      case Dimension::k2D: {
        *output_length = src.size_2d.input_width * src.size_2d.input_height * 2;
        return true;
      }
      case Dimension::k3D: {
        assert_always();
      } break;
      case Dimension::kCube: {
        *output_length =
            src.size_cube.input_width * src.size_cube.input_height * 2 * 6;
        return true;
      }
    }
    return false;
  } else {
    *output_length = src.input_length;
    return true;
  }
}

void TextureCache::WritebackTexture(Texture* texture) {
  VkResult status = VK_SUCCESS;
  VkFence fence = wb_command_pool_->BeginBatch();
  auto alloc = wb_staging_buffer_.Acquire(
      texture->base_region->allocation_info.size, fence);
  if (!alloc) {
    wb_command_pool_->EndBatch();
    return;
  }

  auto command_buffer = wb_command_pool_->AcquireEntry();

  VkCommandBufferBeginInfo begin_info = {
      VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
      nullptr,
      VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
      nullptr,
  };
  vkBeginCommandBuffer(command_buffer, &begin_info);

  // TODO: Transition the texture to a transfer source.

  VkBufferImageCopy region = {
      alloc->offset,
      0,
      0,
      {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1},
      {0, 0, 0},
      {texture->texture_info.width + 1, texture->texture_info.height + 1, 1},
  };

  vkCmdCopyImageToBuffer(command_buffer, texture->base_region->image,
                         VK_IMAGE_LAYOUT_GENERAL,
                         wb_staging_buffer_.gpu_buffer(), 1, &region);

  // TODO: Transition the texture back to a shader resource.

  vkEndCommandBuffer(command_buffer);

  // Submit the command buffer.
  // Submit commands and wait.
  {
    std::lock_guard<std::mutex>(device_->primary_queue_mutex());
    VkSubmitInfo submit_info = {
        VK_STRUCTURE_TYPE_SUBMIT_INFO,
        nullptr,
        0,
        nullptr,
        nullptr,
        1,
        &command_buffer,
        0,
        nullptr,
    };
    status = vkQueueSubmit(device_->primary_queue(), 1, &submit_info, fence);
    CheckResult(status, "vkQueueSubmit");

    if (status == VK_SUCCESS) {
      status = vkQueueWaitIdle(device_->primary_queue());
      CheckResult(status, "vkQueueWaitIdle");
    }
  }

  wb_command_pool_->EndBatch();

  auto dest = memory_->TranslatePhysical(texture->texture_info.guest_address);
  if (status == VK_SUCCESS) {
    std::memcpy(dest, alloc->host_ptr, texture->texture_info.input_length);
  }

  wb_staging_buffer_.Scavenge();
}

void TextureCache::HashTextureBindings(
    XXH64_state_t* hash_state, uint32_t& fetch_mask,
    const std::vector<Shader::TextureBinding>& bindings) {
  for (auto& binding : bindings) {
    uint32_t fetch_bit = 1 << binding.fetch_constant;
    if (fetch_mask & fetch_bit) {
      // We've covered this binding.
      continue;
    }
    fetch_mask |= fetch_bit;

    auto& regs = *register_file_;
    int r = XE_GPU_REG_SHADER_CONSTANT_FETCH_00_0 + binding.fetch_constant * 6;
    auto group =
        reinterpret_cast<const xenos::xe_gpu_fetch_group_t*>(&regs.values[r]);
    auto& fetch = group->texture_fetch;

    XXH64_update(hash_state, &fetch, sizeof(fetch));
  }
}

VkDescriptorSet TextureCache::PrepareTextureSet(
    VkCommandBuffer command_buffer, VkCommandBuffer setup_buffer,
    VkFence completion_fence,
    const std::vector<Shader::TextureBinding>& vertex_bindings,
    const std::vector<Shader::TextureBinding>& pixel_bindings) {
  XXH64_state_t hash_state;
  XXH64_reset(&hash_state, 0);

  // (quickly) Generate a hash.
  uint32_t fetch_mask = 0;
  HashTextureBindings(&hash_state, fetch_mask, vertex_bindings);
  HashTextureBindings(&hash_state, fetch_mask, pixel_bindings);
  uint64_t hash = XXH64_digest(&hash_state);
  for (auto it = texture_sets_.find(hash); it != texture_sets_.end(); ++it) {
    // TODO(DrChat): We need to compare the bindings and ensure they're equal.
    return it->second;
  }

  // Clear state.
  auto update_set_info = &update_set_info_;
  update_set_info->has_setup_fetch_mask = 0;
  update_set_info->image_write_count = 0;

  std::memset(update_set_info, 0, sizeof(update_set_info_));

  // Process vertex and pixel shader bindings.
  // This does things lazily and de-dupes fetch constants reused in both
  // shaders.
  bool any_failed = false;
  any_failed =
      !SetupTextureBindings(command_buffer, setup_buffer, completion_fence,
                            update_set_info, vertex_bindings) ||
      any_failed;
  any_failed =
      !SetupTextureBindings(command_buffer, setup_buffer, completion_fence,
                            update_set_info, pixel_bindings) ||
      any_failed;
  if (any_failed) {
    XELOGW("Failed to setup one or more texture bindings");
    // TODO(benvanik): actually bail out here?
  }

  // Open a new batch of descriptor sets (for this frame)
  if (!descriptor_pool_->has_open_batch()) {
    descriptor_pool_->BeginBatch(completion_fence);
  }

  auto descriptor_set =
      descriptor_pool_->AcquireEntry(texture_descriptor_set_layout_);
  if (!descriptor_set) {
    return nullptr;
  }

  for (uint32_t i = 0; i < update_set_info->image_write_count; i++) {
    update_set_info->image_writes[i].dstSet = descriptor_set;
  }

  // Update the descriptor set.
  if (update_set_info->image_write_count > 0) {
    vkUpdateDescriptorSets(*device_, update_set_info->image_write_count,
                           update_set_info->image_writes, 0, nullptr);
  }

  texture_sets_[hash] = descriptor_set;
  return descriptor_set;
}

bool TextureCache::SetupTextureBindings(
    VkCommandBuffer command_buffer, VkCommandBuffer setup_buffer,
    VkFence completion_fence, UpdateSetInfo* update_set_info,
    const std::vector<Shader::TextureBinding>& bindings) {
  bool any_failed = false;
  for (auto& binding : bindings) {
    uint32_t fetch_bit = 1 << binding.fetch_constant;
    if ((update_set_info->has_setup_fetch_mask & fetch_bit) == 0) {
      // Needs setup.
      any_failed =
          !SetupTextureBinding(command_buffer, setup_buffer, completion_fence,
                               update_set_info, binding) ||
          any_failed;
      update_set_info->has_setup_fetch_mask |= fetch_bit;
    }
  }
  return !any_failed;
}

bool TextureCache::SetupTextureBinding(VkCommandBuffer command_buffer,
                                       VkCommandBuffer setup_buffer,
                                       VkFence completion_fence,
                                       UpdateSetInfo* update_set_info,
                                       const Shader::TextureBinding& binding) {
#if FINE_GRAINED_DRAW_SCOPES
  SCOPE_profile_cpu_f("gpu");
#endif  // FINE_GRAINED_DRAW_SCOPES

  auto& regs = *register_file_;
  int r = XE_GPU_REG_SHADER_CONSTANT_FETCH_00_0 + binding.fetch_constant * 6;
  auto group =
      reinterpret_cast<const xenos::xe_gpu_fetch_group_t*>(&regs.values[r]);
  auto& fetch = group->texture_fetch;

  // Disabled?
  // TODO(benvanik): reset sampler.
  if (fetch.type != 0x2) {
    return false;
  }

  TextureInfo texture_info;
  if (!TextureInfo::Prepare(fetch, &texture_info)) {
    XELOGE("Unable to parse texture fetcher info");
    return false;  // invalid texture used
  }
  SamplerInfo sampler_info;
  if (!SamplerInfo::Prepare(fetch, binding.fetch_instr, &sampler_info)) {
    XELOGE("Unable to parse sampler info");
    return false;  // invalid texture used
  }

  // Search via the base format.
  texture_info.texture_format = GetBaseFormat(texture_info.texture_format);

  auto texture_region = DemandRegion(texture_info, command_buffer, setup_buffer,
                                     completion_fence);
  auto sampler = Demand(sampler_info);
  if (texture_region == nullptr || sampler == nullptr) {
    return false;
  }

  uint16_t swizzle = static_cast<uint16_t>(fetch.swizzle);
  auto view = DemandTextureRegionView(texture_region, swizzle);

  auto image_info =
      &update_set_info->image_infos[update_set_info->image_write_count];
  auto image_write =
      &update_set_info->image_writes[update_set_info->image_write_count];
  update_set_info->image_write_count++;

  // Sanity check, we only have 32 binding slots.
  assert(binding.binding_index < 32);

  image_write->sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
  image_write->pNext = nullptr;
  // image_write->dstSet is set later...
  image_write->dstBinding = 0;
  image_write->dstArrayElement = uint32_t(binding.binding_index);
  image_write->descriptorCount = 1;
  image_write->descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  image_write->pImageInfo = image_info;
  image_write->pBufferInfo = nullptr;
  image_write->pTexelBufferView = nullptr;

  image_info->imageView = view->view;
  image_info->imageLayout = texture_region->image_layout;
  image_info->sampler = sampler->sampler;
  texture_region->texture->in_flight_fence = completion_fence;

  return true;
}

void TextureCache::RemoveInvalidatedTextures() {
  // Clean up any invalidated textures.
  invalidated_textures_mutex_.lock();
  std::unordered_set<Texture*>& invalidated_textures = *invalidated_textures_;
  if (invalidated_textures_ == &invalidated_textures_sets_[0]) {
    invalidated_textures_ = &invalidated_textures_sets_[1];
  } else {
    invalidated_textures_ = &invalidated_textures_sets_[0];
  }
  invalidated_textures_mutex_.unlock();

  // Append all invalidated textures to a deletion queue. They will be deleted
  // when all command buffers using them have finished executing.
  if (!invalidated_textures.empty()) {
    for (auto it = invalidated_textures.begin();
         it != invalidated_textures.end(); ++it) {
      pending_delete_textures_.push_back(*it);
      textures_.erase((*it)->texture_info.hash());
    }

    COUNT_profile_set("gpu/texture_cache/textures", textures_.size());
    COUNT_profile_set("gpu/texture_cache/pending_deletes",
                      pending_delete_textures_.size());
    invalidated_textures.clear();
  }
}

void TextureCache::ClearCache() {
  RemoveInvalidatedTextures();
  for (auto it = textures_.begin(); it != textures_.end(); ++it) {
    while (!FreeTexture(it->second)) {
      // Texture still in use. Busy loop.
      xe::threading::MaybeYield();
    }
  }
  textures_.clear();
  COUNT_profile_set("gpu/texture_cache/textures", 0);

  for (auto it = samplers_.begin(); it != samplers_.end(); ++it) {
    vkDestroySampler(*device_, it->second->sampler, nullptr);
    delete it->second;
  }
  samplers_.clear();
}

void TextureCache::Scavenge() {
  SCOPE_profile_cpu_f("gpu");

  // Close any open descriptor pool batches
  if (descriptor_pool_->has_open_batch()) {
    descriptor_pool_->EndBatch();
  }

  // Free unused descriptor sets
  // TODO(DrChat): These sets could persist across frames, we just need a smart
  // way to detect if they're unused and free them.
  texture_sets_.clear();
  descriptor_pool_->Scavenge();
  staging_buffer_.Scavenge();

  // Kill all pending delete textures.
  RemoveInvalidatedTextures();
  if (!pending_delete_textures_.empty()) {
    for (auto it = pending_delete_textures_.begin();
         it != pending_delete_textures_.end();) {
      if (!FreeTexture(*it)) {
        break;
      }

      it = pending_delete_textures_.erase(it);
    }

    COUNT_profile_set("gpu/texture_cache/pending_deletes",
                      pending_delete_textures_.size());
  }
}

}  // namespace vulkan
}  // namespace gpu
}  // namespace xe
