//
// Created by lgv on 11/25/25.
//


#include <algorithm>
#include <cstdio>
#include <chrono>
#include <cmath>
#include "simulate/mujoco/mujoco_viewer/include/mujoco_viewer.h"
#include "simulate/mujoco/mujoco_viewer/include/array_safety.h"
#include "simulate/mujoco/mujoco_viewer/include/glfw_adapter.h"

namespace cmvr {
    namespace mj = ::mujoco;
    namespace mju = ::mujoco::sample_util;

    using Seconds = std::chrono::duration<double>;

    constexpr double kSyncMisalign = 0.1;
    constexpr double kSimRefreshFraction = 0.7;
    constexpr int kPiPMaxGeom = 100000;

    class PiPGlfwAdapter : public mj::GlfwAdapter {
    public:
        explicit PiPGlfwAdapter(MuJocoViewer *owner)
            : owner_(owner) {}

        void SwapBuffers() override {
            if (owner_) {
                owner_->renderPiP();
            }
            mj::GlfwAdapter::SwapBuffers();
        }

    private:
        MuJocoViewer *owner_;
    };


    static const char *CheckDiverged(int disableflags, const mjData *d) {
        if (disableflags & mjDSBL_AUTORESET) {
            for (mjtWarning w: {mjWARN_BADQACC, mjWARN_BADQVEL, mjWARN_BADQPOS}) {
                if (d->warning[w].number > 0) {
                    return mju_warningText(w, d->warning[w].lastinfo);
                }
            }
        }
        return nullptr;
    }


    mjModel *MuJocoViewer::LoadModelSimple(const char *file,
                                           char *load_error,
                                           int error_sz) {
        load_error[0] = '\0';
        if (!file || !file[0]) {
            std::snprintf(load_error, error_sz, "empty model filename");
            return nullptr;
        }

        mjModel *mnew = nullptr;
        const char *dot = std::strrchr(file, '.');

        if (dot && std::strcmp(dot, ".mjb") == 0) {
            mnew = mj_loadModel(file, nullptr);
            if (!mnew) {
                std::snprintf(load_error, error_sz,
                              "could not load binary model: %s", file);
            }
        } else {
            mnew = mj_loadXML(file, nullptr, load_error, error_sz);
            if (load_error[0]) {
                int len = std::strlen(load_error);
                if (len > 0 && load_error[len - 1] == '\n') {
                    load_error[len - 1] = '\0';
                }
            }
        }
        return mnew;
    }

    MuJocoViewer::MuJocoViewer(const char *model_path)
        : model_path_(model_path ? model_path : "") {
        std::printf("MuJoCo version %s\n", mj_versionString());
        if (mjVERSION_HEADER != mj_version()) {
            mju_error("Headers and library have different versions");
        }

        initSim();
    }

    void MuJocoViewer::initSim() {
        mjv_defaultCamera(&cam_);
        mjv_defaultOption(&opt_);
        mjv_defaultPerturb(&pert_);
        mjv_defaultCamera(&pip_cam_);
        mjv_defaultScene(&pip_scene_);

        auto platform_ui = std::make_unique<PiPGlfwAdapter>(this);
        sim_ = std::make_unique<mj::Simulate>(
            std::move(platform_ui),
            &cam_, &opt_, &pert_,
            /*is_passive=*/false
        );
    }

    MuJocoViewer::~MuJocoViewer() {
        if (sim_) {
            sim_->exitrequest.store(1);
        }
        if (physics_started_.load() && physics_thread_.joinable()) {
            physics_thread_.join();
        }

        if (pip_scene_inited_) {
            mjv_freeScene(&pip_scene_);
            pip_scene_inited_ = false;
        }

        if (d_) mj_deleteData(d_);
        if (m_) mj_deleteModel(m_);
    }

    void MuJocoViewer::physicsThreadFunc() {
        // 1) 加载模型
        char error[1024];
        m_ = LoadModelSimple(model_path_.c_str(), error, sizeof(error));
        if (!m_) {
            mju_error("Load model failed: %s", error);
        }

        d_ = mj_makeData(m_);
        if (!d_) {
            mj_deleteModel(m_);
            m_ = nullptr;
            mju_error("mj_makeData failed");
        }

        // 2) 用官方的方式：先告诉 Simulate 要加载，然后等 RenderLoop 处理
        sim_->Load(m_, d_, model_path_.c_str());

        {
            // official main 里也是在 Load 之后先 forward 一下
            const std::unique_lock<std::recursive_mutex> lock(sim_->mtx);
            mj_forward(m_, d_);
        }

        physics_started_.store(true);

        // 3) 进入物理循环
        physicsLoop();
    }

    void MuJocoViewer::getCameraState(double &distance,
                                      double &azimuth,
                                      double &elevation) const {
        distance = cam_.distance;
        azimuth = cam_.azimuth;
        elevation = cam_.elevation;
    }

    void MuJocoViewer::setupCamera(double distance,
                                   double azimuth,
                                   double elevation) {
        // 相机初始化：根据模型设置一个合理的 free camera
        if (m_) {
            mjv_defaultFreeCamera(m_, &cam_);
            // 以模型中心为观察点
            cam_.lookat[0] = m_->stat.center[0];
            cam_.lookat[1] = m_->stat.center[1];
            cam_.lookat[2] = m_->stat.center[2];
        } else {
            mjv_defaultCamera(&cam_);
        }

        cam_.type = mjCAMERA_FREE;
        cam_.distance = distance; // 距离
        cam_.azimuth = azimuth; // 方位角（度）
        cam_.elevation = elevation; // 俯仰角（度）
    }

    void MuJocoViewer::enablePiPCamera(const char *camera_name,
                                       int left,
                                       int bottom,
                                       int width,
                                       int height) {
        pip_enabled_ = true;
        pip_camera_name_ = camera_name ? camera_name : "";
        pip_camera_id_ = -1;
        pip_left_ = left;
        pip_bottom_ = bottom;
        pip_width_ = width > 0 ? width : 320;
        pip_height_ = height > 0 ? height : 240;
        pip_custom_pos_ = true;
    }

    void MuJocoViewer::disablePiPCamera() {
        pip_enabled_ = false;
    }



    void MuJocoViewer::renderPiP() {
        if (!pip_enabled_ || !sim_ || !m_ || !d_) return;
        if (pip_camera_name_.empty()) return;

        if (pip_camera_id_ < 0) {
            pip_camera_id_ = mj_name2id(m_, mjOBJ_CAMERA, pip_camera_name_.c_str());
            if (pip_camera_id_ < 0) {
                return;
            }
        }

        auto [fb_width, fb_height] = sim_->platform_ui->GetFramebufferSize();
        if (fb_width <= 0 || fb_height <= 0) return;

        int left = 0;
        int bottom = 0;
        int width = 0;
        int height = 0;

        if (pip_custom_pos_) {
            left = std::max(0, std::min(pip_left_, fb_width - 1));
            bottom = std::max(0, std::min(pip_bottom_, fb_height - 1));
            width = std::min(pip_width_, fb_width - left);
            height = std::min(pip_height_, fb_height - bottom);
        } else {
            width = std::min(pip_width_, fb_width - 2 * pip_margin_);
            height = std::min(pip_height_, fb_height - 2 * pip_margin_);
            left = fb_width - pip_margin_ - width;
            bottom = pip_margin_;
        }

        if (width <= 0 || height <= 0) return;

        mjrRect rect;
        rect.width = width;
        rect.height = height;
        rect.left = left;
        rect.bottom = bottom;

        const std::unique_lock<std::recursive_mutex> lock(sim_->mtx);

        if (!pip_scene_inited_ || pip_scene_model_ != m_) {
            if (pip_scene_inited_) {
                mjv_freeScene(&pip_scene_);
            }
            mjv_makeScene(m_, &pip_scene_, kPiPMaxGeom);
            pip_scene_inited_ = true;
            pip_scene_model_ = m_;
        }

        pip_cam_.type = mjCAMERA_FIXED;
        pip_cam_.fixedcamid = pip_camera_id_;
        pip_cam_.trackbodyid = -1;

        mjv_updateScene(m_, d_, &opt_, &pert_, &pip_cam_, mjCAT_ALL, &pip_scene_);
        mjr_render(rect, &pip_scene_, &sim_->platform_ui->mjr_context());

        {
            std::lock_guard<std::mutex> lock(pip_rgb_mtx_);
            const int w = rect.width;
            const int h = rect.height;
            if (w > 0 && h > 0) {
                pip_rgb_.resize(static_cast<size_t>(3 * w * h));
                pip_depth_.resize(static_cast<size_t>(w * h));

                // 同时读 RGB 和 depth(z-buffer 0..1)
                mjr_readPixels(pip_rgb_.data(), pip_depth_.data(),
                               rect, &sim_->platform_ui->mjr_context());

                // OpenGL 像素原点在左下，需要竖直翻转 RGB 和 depth
                for (int r = 0; r < h / 2; ++r) {
                    // flip rgb row
                    unsigned char *top_row = pip_rgb_.data() + 3 * w * r;
                    unsigned char *bottom_row = pip_rgb_.data() + 3 * w * (h - 1 - r);
                    std::swap_ranges(top_row, top_row + 3 * w, bottom_row);

                    // flip depth row
                    float *top_d = pip_depth_.data() + w * r;
                    float *bot_d = pip_depth_.data() + w * (h - 1 - r);
                    std::swap_ranges(top_d, top_d + w, bot_d);
                }

                pip_rgb_width_ = w;
                pip_rgb_height_ = h;
                pip_rgb_valid_ = true;
                ++pip_frame_id_;   // 新帧
            }
        }
    }

    void MuJocoViewer::physicsLoop() {
        using Clock = mj::Simulate::Clock;

        std::chrono::time_point<Clock> syncCPU;
        mjtNum syncSim = 0;

        while (!sim_->exitrequest.load()) {
            // 给 UI 线程一点时间
            if (sim_->run && sim_->busywait) {
                std::this_thread::yield();
            } else {
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
            }

            {
                const std::unique_lock<std::recursive_mutex> lock(sim_->mtx);
                if (!m_ || !d_) continue;

                if (sim_->run) {
                    bool stepped = false;
                    const auto startCPU = Clock::now();
                    const auto elapsedCPU = startCPU - syncCPU;
                    double elapsedSim = d_->time - syncSim;

                    double slowdown = 100.0 / sim_->percentRealTime[sim_->real_time_index];
                    bool misaligned =
                            std::abs(Seconds(elapsedCPU).count() / slowdown - elapsedSim) > kSyncMisalign;

                    if (elapsedSim < 0 || elapsedCPU.count() < 0 ||
                        syncCPU.time_since_epoch().count() == 0 ||
                        misaligned || sim_->speed_changed) {
                        syncCPU = startCPU;
                        syncSim = d_->time;
                        sim_->speed_changed = false;

                        sim_->InjectNoise(sim_->key);


                        controlDispatch(m_, d_);

                        mj_step(m_, d_);
                        const char *msg = CheckDiverged(m_->opt.disableflags, d_);
                        if (msg) {
                            sim_->run = 0;
                            mju::strcpy_arr(sim_->load_error, msg);
                        } else {
                            stepped = true;
                        }
                    } else {
                        bool measured = false;
                        mjtNum prevSim = d_->time;
                        double refreshTime = kSimRefreshFraction / sim_->refresh_rate;

                        while (Seconds((d_->time - syncSim) * slowdown) <
                               (Clock::now() - syncCPU) &&
                               (Clock::now() - startCPU) < Seconds(refreshTime)) {
                            if (!measured && elapsedSim) {
                                sim_->measured_slowdown =
                                        Seconds(elapsedCPU).count() / elapsedSim;
                                measured = true;
                            }

                            sim_->InjectNoise(sim_->key);


                            controlDispatch(m_, d_);

                            mj_step(m_, d_);
                            const char *msg = CheckDiverged(m_->opt.disableflags, d_);
                            if (msg) {
                                sim_->run = 0;
                                mju::strcpy_arr(sim_->load_error, msg);
                            } else {
                                stepped = true;
                            }

                            if (d_->time < prevSim) {
                                break;
                            }
                        }
                    }

                    if (stepped) {
                        sim_->AddToHistory();
                    }
                } else {
                    mj_forward(m_, d_);
                    if (sim_->pause_update) {
                        mju_copy(d_->qacc_warmstart, d_->qacc, m_->nv);
                    }
                    sim_->speed_changed = true;
                }

                if (d_->time == 0.0) {
                    onReset(m_, d_);
                }
            }
        }
    }

    void MuJocoViewer::controlDispatch(mjModel *m, mjData *d) {
        // 第一次：只跑一次 initOnce
        if (!inited_) {
            initOnce(m, d);
            inited_ = true;
            last_time_ = d->time;
        }

        // 检测 reset：时间从大跳到小（比如从 3.2 回到 0）
        if (d->time < last_time_) {
            onReset(m, d);
            inited_ = false;
        }
        last_time_ = d->time;

        // 每步控制
        controlCallback(m, d);
    }

    void MuJocoViewer::run() {
        // 先起 physics 线程（里面会 sim_->Load）
        physics_thread_ = std::thread(&MuJocoViewer::physicsThreadFunc, this);

        // 当前线程跑 UI + 渲染（阻塞）
        sim_->RenderLoop();

        // UI 退出后，通知 physics 线程结束
        sim_->exitrequest.store(2);

        if (physics_thread_.joinable()) {
            physics_thread_.join();
        }
    }


    uint64_t MuJocoViewer::getPiPCameraFrameId() const {
        std::lock_guard<std::mutex> lock(pip_rgb_mtx_);
        return pip_frame_id_;
    }

    bool MuJocoViewer::getPiPCameraRGBD(std::vector<unsigned char> &rgb,
                                        std::vector<float> &depth,
                                        int &width,
                                        int &height,
                                        uint64_t &frame_id) const {
        std::lock_guard<std::mutex> lock(pip_rgb_mtx_);
        if (!pip_rgb_valid_ || pip_rgb_.empty()) return false;

        rgb = pip_rgb_;
        depth = pip_depth_;
        width = pip_rgb_width_;
        height = pip_rgb_height_;
        frame_id = pip_frame_id_;
        return true;
    }
} // namespace cmvr