// Copyright 2026 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
//
//     https://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.

// The web viewer's headless Studio UI.
//
// This pybind11 module provides a HeadlessUi class that:
// 1. Creates a headless ImGui context (no window, no renderer).
// 2. Connects as a netimgui client, streaming ImGui draw data to a remote
//    viewer (the browser's web_client, bridged through web_server.py).
// 3. Receives input events from the remote viewer and injects them into
//    the ImGui context.

#include <chrono>
#include <cstdint>
#include <string>
#include <string_view>
#include <thread>
#include <vector>

#include <imgui.h>
#include <implot.h>
#include <mujoco/experimental/platform/ux/fonts.h>
#include <NetImgui_Api.h>
#include "google/logging.h"
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

namespace py = pybind11;

// Headless ImGui + netimgui viewer for the MuJoCo web viewer.
//
// This class manages a headless ImGui context that streams its draw data
// via the netimgui protocol to a remote viewer. Unlike the NativeViewer's
// Viewer class (native_viewer.cc), this does NOT create a window, initialize
// a renderer, or handle mouse/keyboard input directly. All rendering and
// input handling happens on the remote client (web_client.cc in the browser).
//
// The lifecycle follows the SampleNoBackend pattern:
//   Client_Startup()   — create context, load fonts, init NetImgui
//   Client_Connect()   — manage connection state (called each frame)
//   Client_Shutdown()  — release resources

// Initialize the Dear ImGui Context and the NetImgui library.
// Based on SampleNoBackend::Client_Startup() from
// netimgui/Code/Sample/SampleNoBackend/SampleNoBackend.cpp
static bool Client_Startup(ImGuiContext*& context,
                           const std::string& assets_dir) {
  IMGUI_CHECKVERSION();
  context = ImGui::CreateContext();
  ImGui::SetCurrentContext(context);
  ImPlot::CreateContext();

  ImGuiIO& io = ImGui::GetIO();
  io.ConfigFlags |= ImGuiConfigFlags_DockingEnable;
  io.BackendFlags |= ImGuiBackendFlags_HasGamepad;
  io.IniFilename = nullptr;
  io.ConfigDpiScaleFonts = true;
  io.ConfigDpiScaleViewports = true;
  io.DisplaySize = ImVec2(1400, 720);

  // Disable ImGui auto-repeat in headless mode since the browser handles key
  // repeat at the OS level, avoiding the frame amplification issue where many
  // frames with identical key-down states trigger unwanted repeats.
  io.KeyRepeatDelay = 9999.0f;

  // Initialize the main viewport's DPI scale. Without this, the headless
  // context leaves DpiScale at 0.0f (no platform backend sets it), which
  // triggers an assertion in ImGui::SetCurrentViewport() when
  // BeginMainMenuBar() is called.
  ImGuiViewport* main_vp = ImGui::GetMainViewport();
  if (main_vp) {
    main_vp->DpiScale = 1.0f;
  }

  ImGui::StyleColorsLight();

  // The Studio font set shared with the native viewer and the browser client.
  mujoco::platform::AddStudioFonts([&assets_dir](std::string_view filename) {
    return mujoco::platform::LoadFontAsset(assets_dir, filename);
  });

  if (!NetImgui::Startup()) {
    LOG(Error, "NetImgui::Startup() failed");
    return false;
  }

  return true;
}

// Release resources.
// Based on SampleNoBackend::Client_Shutdown() from
// netimgui/Code/Sample/SampleNoBackend/SampleNoBackend.cpp
static void Client_Shutdown(ImGuiContext*& context) {
  NetImgui::Shutdown();
  ImPlot::DestroyContext();
  if (context) {
    ImGui::DestroyContext(context);
    context = nullptr;
  }
}

// Manage connection to the netimgui proxy.
// Based on SampleNoBackend::Client_Connect() from
// netimgui/Code/Sample/SampleNoBackend/SampleNoBackend.cpp
static void Client_Connect(const char* title, int port) {
  bool connected = NetImgui::IsConnected();
  bool pending = NetImgui::IsConnectionPending();

  if (!connected && !pending) {
    static std::chrono::steady_clock::time_point last_reconnect_time =
        std::chrono::steady_clock::now();
    const std::chrono::steady_clock::time_point now =
        std::chrono::steady_clock::now();
    if (now - last_reconnect_time > std::chrono::seconds(1)) {
      last_reconnect_time = now;
      VLOG(1, "Retrying ConnectToApp...");
      NetImgui::ConnectToApp(title, "127.0.0.1", port);
    }
  }

  static bool last_connected = false;
  if (connected != last_connected) {
    VLOG(1, "Status change: Connected=%s", connected ? "true" : "false");
    last_connected = connected;
  }
}

class HeadlessUi {
 public:
  HeadlessUi(const std::string& title, int port, const std::string& assets_dir)
      : title_(title), port_(port) {
    if (!Client_Startup(context_, assets_dir)) {
      return;
    }

    VLOG(1, "Calling ConnectToApp('%s', '127.0.0.1', %d)", title_.c_str(),
         port_);
    bool connect_result =
        NetImgui::ConnectToApp(title_.c_str(), "127.0.0.1", port_);
    VLOG(1, "ConnectToApp returned: %s", connect_result ? "true" : "false");
    VLOG(1, "IsConnected: %s, IsConnectionPending: %s",
         NetImgui::IsConnected() ? "true" : "false",
         NetImgui::IsConnectionPending() ? "true" : "false");
  }

  ~HeadlessUi() { Client_Shutdown(context_); }

  bool NewFrame() {
    py::gil_scoped_release no_gil;
    ImGui::SetCurrentContext(context_);
    static int frame_count = 0;
    frame_count++;

    // Returns true with an active ImGui frame once a browser is connected and
    // ready. Returns false (no frame) while no browser is connected so the
    // caller's loop keeps running and can drain messages (e.g., to respond to
    // an ExitEvent for shutdown).
    std::chrono::steady_clock::time_point last_signal_check =
        std::chrono::steady_clock::now();
    while (true) {
      // While no browser is connected, periodically check for
      // Python signals so Ctrl+C interrupts the wait instead of hanging.
      const std::chrono::steady_clock::time_point now =
          std::chrono::steady_clock::now();
      if (now - last_signal_check > std::chrono::milliseconds(200)) {
        last_signal_check = now;
        py::gil_scoped_acquire acquire;
        if (PyErr_CheckSignals() != 0) {
          throw py::error_already_set();
        }
      }

      Client_Connect(title_.c_str(), port_);

      if (!NetImgui::IsConnected()) {
        // Not connected; return false so the caller's loop is not blocked.
        is_drawing_remote_ = false;
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
        return false;
      }

      bool new_frame_result = NetImgui::NewFrame(false);
      is_drawing_remote_ = new_frame_result;
      if (!new_frame_result) {
        // Connected but NetImgui not ready for a draw — wait and retry.
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
        continue;
      }

      break;
    }

    VLOG(1, "Frame %d: DrawingRemote=%s", frame_count,
         is_drawing_remote_ ? "Y" : "N");
    return true;
  }

  void EndFrame() {
    py::gil_scoped_release no_gil;
    ImGui::SetCurrentContext(context_);
    static int end_frame_count = 0;
    end_frame_count++;
    if (!is_drawing_remote_) {
      // No frame was started — nothing to end.
      return;
    }
    VLOG(1, "EndFrame %d: sending remote draw data", end_frame_count);
    NetImgui::EndFrame();
  }

  // Uploads an RGB/RGBA image to the browser over the NetImgui texture
  // channel so handlers can display it with imgui.Image(). Returns the
  // texture id to use (allocates one when tex_id == 0). This is the UI-link
  // counterpart of the native renderer's UploadImage.
  uintptr_t UploadImage(uintptr_t tex_id, const py::bytes& pixels, int width,
                        int height, int bpp) {
    if (tex_id == 0) {
      tex_id = next_tex_id_++;
    }
    std::string data = pixels;
    if (width <= 0 || height <= 0 ||
        data.size() < static_cast<size_t>(width) * height * bpp) {
      LOG(Error, "UploadImage: bad dimensions %dx%dx%d for %zu bytes", width,
          height, bpp, data.size());
      return tex_id;
    }

    // NetImgui transfers RGBA8; expand RGB if needed.
    std::vector<uint8_t> rgba;
    const void* upload_data = data.data();
    if (bpp == 3) {
      rgba.resize(static_cast<size_t>(width) * height * 4);
      for (size_t p = 0; p < static_cast<size_t>(width) * height; ++p) {
        rgba[p * 4 + 0] = data[p * 3 + 0];
        rgba[p * 4 + 1] = data[p * 3 + 1];
        rgba[p * 4 + 2] = data[p * 3 + 2];
        rgba[p * 4 + 3] = 255;
      }
      upload_data = rgba.data();
    } else if (bpp != 4) {
      LOG(Error, "UploadImage: unsupported bpp %d (expected 3 or 4)", bpp);
      return tex_id;
    }

    py::gil_scoped_release no_gil;
    ImGui::SetCurrentContext(context_);
    NetImgui::SendDataTexture(
        static_cast<ImTextureID>(tex_id), const_cast<void*>(upload_data),
        static_cast<uint16_t>(width), static_cast<uint16_t>(height),
        NetImgui::eTexFormat::kTexFmtRGBA8);
    return tex_id;
  }

  intptr_t GetContext() const { return reinterpret_cast<intptr_t>(context_); }

  // The ImPlot context created alongside the ImGui context. Python must pass
  // this to ux.set_implot_context: extension modules each hold their own copy
  // of the ImPlot globals, so the context pointer has to be shared explicitly
  // (same pattern as get_context/set_imgui_context).
  intptr_t GetImPlotContext() const {
    return reinterpret_cast<intptr_t>(ImPlot::GetCurrentContext());
  }

 private:
  std::string title_;
  int port_;
  ImGuiContext* context_ = nullptr;
  bool is_drawing_remote_ = false;
  // User texture ids start well above the ids ImGui's managed texture system
  // (font atlas) hands out, so the two can never collide in the browser's
  // texture map.
  uintptr_t next_tex_id_ = 0x10000;
};

PYBIND11_MODULE(headless_ui, m, pybind11::mod_gil_not_used()) {
  m.doc() = "MuJoCo web viewer headless Studio UI, streamed via NetImgui";

  py::class_<HeadlessUi>(m, "HeadlessUi")
      .def(py::init<const std::string&, int, const std::string&>(),
           py::arg("title"), py::arg("port") = 8888, py::arg("assets_dir") = "")
      .def("new_frame", &HeadlessUi::NewFrame,
           "Starts a headless ImGui frame, returning True once the frame is "
           "active. When a browser is viewing the page (via the URL printed "
           "at startup), this returns at the browser's requested frame rate. "
           "When no browser is viewing, it returns False (no frame) after a "
           "short wait, so the caller's loop keeps running and can shut down.")
      .def("end_frame", &HeadlessUi::EndFrame)
      .def("get_context", &HeadlessUi::GetContext)
      .def("get_implot_context", &HeadlessUi::GetImPlotContext)
      .def("upload_image", &HeadlessUi::UploadImage);
}