// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "experimental/platform/ux/imgui_bridge.h"

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <memory>
#include <utility>
#include <vector>

#include <imgui.h>
#include <math/mat3.h>
#include <math/vec3.h>
#include <mujoco/mujoco.h>
#include "render/filament/mjrfilament.h"
#include "render/filament/mjrfilament_cpp.h"

namespace mujoco {

using filament::math::float3;
using filament::math::mat3f;

ImguiBridge::ImguiBridge(mjrfContext* ctx) : ctx_(ctx) {
  mjrfSceneParams params;
  mjrf_defaultSceneParams(&params);
  scene_ = CreateScene(ctx_, params);
}

ImguiBridge::~ImguiBridge() {
  PrepareRenderables(0);

  // Destroy all textures tracked by ImGui.
  if (ImGui::GetCurrentContext()) {
    for (ImTextureData* tex : ImGui::GetPlatformIO().Textures) {
      if (tex->Status != ImTextureStatus_Destroyed) {
        DestroyTexture(tex);
      }
    }
  }
}

uintptr_t ImguiBridge::UploadImage(uintptr_t tex_id, const uint8_t* pixels,
                                   int width, int height, int bpp) {
  if (bpp != 4 && bpp != 3) {
    mju_error("Unsupported image bpp. Got %d, wanted 3 or 4", bpp);
  }

  if (pixels == nullptr) {
    // If the pixels are nullptr, we destroy the texture.
    if (tex_id != 0) {
      textures_.erase(tex_id);
    }
    return 0;
  }

  // Assign a new texture ID.
  if (tex_id == 0) {
    tex_id = next_tex_id_++;
  }

  mjrfTexture* texture = GetTexture(tex_id);

  // If the texture does not exist or the dimensions have changed, we create a
  // new texture.
  if (texture == nullptr || mjrf_getTextureWidth(texture) != width ||
      mjrf_getTextureHeight(texture) != height) {
    mjrfTextureConfig config;
    mjrf_defaultTextureConfig(&config);
    config.width = width;
    config.height = height;
    config.sampler_type = mjTEXTURE_2D;
    config.format = bpp == 4 ? mjPIXEL_FORMAT_RGBA8 : mjPIXEL_FORMAT_RGB8;
    config.color_space = mjCOLORSPACE_LINEAR;
    UniquePtr<mjrfTexture> new_texture = ::mujoco::CreateTexture(ctx_, config);
    texture = new_texture.get();
    textures_.insert_or_assign(tex_id, std::move(new_texture));
  }

  // Create a copy of the image to pass it to filament as we don't know the
  // lifetime of the data.
  const size_t num_bytes = width * height * bpp;
  std::byte* bytes = new std::byte[num_bytes];
  const auto callback =
      +[](void* user) { delete[] reinterpret_cast<std::byte*>(user); };

  mjrfTextureData texture_data;
  mjrf_defaultTextureData(&texture_data);
  texture_data.bytes = bytes;
  texture_data.num_bytes = num_bytes;
  texture_data.user_data = bytes;
  texture_data.release = callback;

  std::memcpy(bytes, pixels, num_bytes);
  mjrf_setTextureData(texture, &texture_data);
  return tex_id;
}

void ImguiBridge::CreateTexture(ImTextureData* data) {
  if (data->Format != ImTextureFormat_RGBA32) {
    mju_error("Unsupported texture format.");
  }

  mjrfTextureConfig config;
  mjrf_defaultTextureConfig(&config);
  config.width = data->Width;
  config.height = data->Height;
  config.sampler_type = mjTEXTURE_2D;
  config.format = mjPIXEL_FORMAT_RGBA8;
  config.color_space = mjCOLORSPACE_LINEAR;

  const uintptr_t tex_id = next_tex_id_++;
  textures_.insert_or_assign(tex_id, ::mujoco::CreateTexture(ctx_, config));
  data->SetTexID((ImTextureID)tex_id);
  UpdateTexture(data);
}

void ImguiBridge::UpdateTexture(ImTextureData* data) {
  auto iter = textures_.find(data->TexID);
  if (iter == textures_.end()) {
    mju_error("Texture not found: %llu", data->TexID);
  }

  mjrfTextureData texture_data;
  mjrf_defaultTextureData(&texture_data);
  texture_data.bytes = data->GetPixels();
  texture_data.num_bytes = data->Width * data->Height * 4;
  texture_data.user_data = nullptr;
  texture_data.release = nullptr;
  mjrf_setTextureData(iter->second.get(), &texture_data);
  data->SetStatus(ImTextureStatus_OK);
}

void ImguiBridge::DestroyTexture(ImTextureData* data) {
  auto iter = textures_.find(data->TexID);
  if (iter != textures_.end()) {
    textures_.erase(data->TexID);
    data->SetTexID(ImTextureID_Invalid);
    data->SetStatus(ImTextureStatus_Destroyed);
  }
}

mjrfTexture* ImguiBridge::GetTexture(uintptr_t tex_id) const {
  auto iter = textures_.find(tex_id);
  if (iter == textures_.end()) {
    return nullptr;
  }
  return iter->second.get();
}

void ImguiBridge::Update() {
  if (!ImGui::GetCurrentContext()) {
    PrepareRenderables(0);
    return;
  }

  // Prepare the imgui draw commands. We must call this function even if we do
  // not plan on rendering anything to ensure imgui state is updated.
  ImGui::Render();

  ImGuiIO& io = ImGui::GetIO();
  const ImVec2& size = io.DisplaySize;
  const ImVec2& scale = io.DisplayFramebufferScale;
  ImDrawData* commands = ImGui::GetDrawData();
  if (!commands || size.x == 0 || size.y == 0) {
    PrepareRenderables(0);
    return;
  }
  commands->ScaleClipRects(scale);

  // 2 floats for position, 2 floats for uv, 4 bytes for color.
  constexpr size_t kExpectedVertexSize =
      sizeof(float) * 4 + sizeof(uint8_t) * 4;

  int num_elements = 0;
  for (int n = 0; n < commands->CmdListsCount; ++n) {
    const ImDrawList* cmds = commands->CmdLists[n];
    if (kExpectedVertexSize != sizeof(cmds->VtxBuffer.Data[0])) {
      mju_error("Invalid vertex buffer size.");
    }
    if (sizeof(uint16_t) != sizeof(cmds->IdxBuffer.Data[0])) {
      mju_error("Invalid index buffer size.");
    }
    num_elements += cmds->CmdBuffer.size();
  }

  if (commands->Textures != nullptr) {
    for (ImTextureData* tex : *commands->Textures) {
      if (tex->Status == ImTextureStatus_WantCreate) {
        CreateTexture(tex);
      } else if (tex->Status == ImTextureStatus_WantUpdates) {
        UpdateTexture(tex);
      } else if (tex->Status == ImTextureStatus_WantDestroy &&
                 tex->UnusedFrames > 0) {
        DestroyTexture(tex);
      }
    }
  }

  PrepareRenderables(num_elements);
  if (num_elements == 0) {
    return;
  }

  meshes_.clear();
  int renderable_index = 0;
  for (int n = 0; n < commands->CmdListsCount; ++n) {
    const ImDrawList* cmds = commands->CmdLists[n];

    mjrfMeshData data;
    mjrf_defaultMeshData(&data);
    data.nattributes = 3;
    data.attributes[0].usage = mjVERTEX_ATTRIBUTE_USAGE_POSITION;
    data.attributes[0].type = mjVERTEX_ATTRIBUTE_TYPE_FLOAT2;
    data.attributes[0].bytes = cmds->VtxBuffer.Data;
    data.attributes[1].usage = mjVERTEX_ATTRIBUTE_USAGE_UV;
    data.attributes[1].type = mjVERTEX_ATTRIBUTE_TYPE_FLOAT2;
    data.attributes[1].bytes = cmds->VtxBuffer.Data + sizeof(float) * 2;
    data.attributes[2].usage = mjVERTEX_ATTRIBUTE_USAGE_COLOR;
    data.attributes[2].type = mjVERTEX_ATTRIBUTE_TYPE_UBYTE4;
    data.attributes[2].bytes = cmds->VtxBuffer.Data + sizeof(float) * 4;
    data.interleaved = true;
    data.num_vertices = cmds->VtxBuffer.Size;
    data.num_indices = cmds->IdxBuffer.Size;
    data.indices = cmds->IdxBuffer.Data;
    data.index_type = mjINDEX_TYPE_U16;
    data.primitive_type = mjMESH_PRIMITIVE_TYPE_TRIANGLES;
    meshes_.push_back(CreateMesh(ctx_, data));

    const mjrfMesh* mesh = meshes_.back().get();

    int index_offset = 0;
    for (const ImDrawCmd& command : cmds->CmdBuffer) {
      const int width = size.x * scale.x;
      const int height = size.y * scale.y;

      UniquePtr<mjrfRenderable>& renderable = renderables_[renderable_index];
      mjrf_setRenderableMesh(renderable.get(), mesh, index_offset,
                             command.ElemCount);

      mjrfMaterial material;
      mjrf_defaultMaterial(&material);
      material.color_texture = GetTexture(command.GetTexID());

      material.decor_ux = true;
      material.scissor[0] = command.ClipRect.x;
      material.scissor[1] = height - command.ClipRect.w;
      material.scissor[2] = command.ClipRect.z - command.ClipRect.x;
      material.scissor[3] = command.ClipRect.w - command.ClipRect.y;
      // Modal dialogs try to cover the whole window, but also a little outside
      // of it. This doesn't work well with filament's scissor test, so we clip
      // them to the window.
      if (material.scissor[0] < 0 || material.scissor[1] < 0) {
        material.scissor[0] = 0;
        material.scissor[1] = 0;
        material.scissor[2] = width;
        material.scissor[3] = height;
      }
      mjrf_setRenderableMaterial(renderable.get(), &material);

      const float size[] = {scale.x, scale.y, 1.0f};
      mjrf_setRenderableSize(renderable.get(), size);

      index_offset += command.ElemCount;
      ++renderable_index;
    }
  }
}

void ImguiBridge::PrepareRenderables(int count) {
  while (renderables_.size() < count) {
    mjrfRenderableParams params;
    mjrf_defaultRenderableParams(&params);
    params.cast_shadows = false;
    params.receive_shadows = false;
    params.blend_order = static_cast<std::uint16_t>(renderables_.size() + 1);
    auto& renderable = renderables_.emplace_back(CreateRenderable(ctx_, params));
    mjrf_addRenderableToScene(scene_.get(), renderable.get());
  }
  while (renderables_.size() > count) {
    mjrf_removeRenderableFromScene(scene_.get(), renderables_.back().get());
    renderables_.pop_back();
  }
}

mjrfScene* ImguiBridge::GetScene() const { return scene_.get(); }

mjrCamera ImguiBridge::GetCamera(int width, int height) const {
  mjrCamera camera;
  camera.orthographic = true;
  camera.pos[0] = 0.0f;
  camera.pos[1] = 0.0f;
  camera.pos[2] = 1.0f;
  camera.forward[0] = 0.0f;
  camera.forward[1] = 0.0f;
  camera.forward[2] = -1.0f;
  camera.up[0] = 0.0f;
  camera.up[1] = 1.0f;
  camera.up[2] = 0.0f;
  camera.frustum_top = 0.0f;
  camera.frustum_near = 0.0f;
  camera.frustum_far = 1.0f;
  camera.frustum_center = width / 2.0f;
  camera.frustum_width = width / 2.0f;
  camera.frustum_bottom = height;
  return camera;
}

}  // namespace mujoco