#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <cstdint>
#include <cstring>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <string>

#ifdef DJI_SDK_AVAILABLE
#include "dji_ack.hpp"
#include "dji_broadcast.hpp"
#include "dji_control.hpp"
#include "dji_status.hpp"
#include "dji_telemetry.hpp"
#include "dji_vehicle.hpp"
#endif

namespace py = pybind11;

namespace {

py::dict make_attitude(double q0, double q1, double q2, double q3) {
    py::dict d;
    d["q0"] = q0;
    d["q1"] = q1;
    d["q2"] = q2;
    d["q3"] = q3;
    return d;
}

py::dict make_position(double lat, double lon, double alt) {
    py::dict d;
    d["lat"] = lat;
    d["lon"] = lon;
    d["alt"] = alt;
    return d;
}

py::dict make_velocity(double vx, double vy, double vz) {
    py::dict d;
    d["vx"] = vx;
    d["vy"] = vy;
    d["vz"] = vz;
    return d;
}

py::dict make_rc(int pitch, int roll, int yaw, int throttle) {
    py::dict d;
    d["pitch"] = pitch;
    d["roll"] = roll;
    d["yaw"] = yaw;
    d["throttle"] = throttle;
    return d;
}

#ifdef DJI_SDK_AVAILABLE

namespace osdk = DJI::OSDK;

constexpr int kDefaultWaitTimeoutMs = 1000;

std::mutex g_vehicle_mutex;
std::unique_ptr<osdk::Vehicle> g_vehicle;

osdk::Vehicle* require_vehicle() {
    std::lock_guard<std::mutex> lock(g_vehicle_mutex);
    if (!g_vehicle) {
        throw std::runtime_error("DJI vehicle is not connected");
    }
    return g_vehicle.get();
}

bool connect_real(const std::string& device, std::uint32_t baudrate, bool enable_advanced_sensing, bool enable_broadcast) {
    std::lock_guard<std::mutex> lock(g_vehicle_mutex);
    auto vehicle = std::make_unique<osdk::Vehicle>(device.c_str(), baudrate, enable_advanced_sensing, enable_broadcast);
    if (!vehicle) {
        return false;
    }
    if (!vehicle->protocolLayer || !vehicle->control) {
        return false;
    }
    g_vehicle = std::move(vehicle);
    return true;
}

void disconnect_real() {
    std::lock_guard<std::mutex> lock(g_vehicle_mutex);
    g_vehicle.reset();
}

bool is_connected_real() {
    std::lock_guard<std::mutex> lock(g_vehicle_mutex);
    return static_cast<bool>(g_vehicle);
}

bool activate_real(py::object app_id_obj, py::object app_key_obj, py::object api_level_obj) {
    auto* vehicle = require_vehicle();

    osdk::Vehicle::ActivateData data{};
    if (!app_id_obj.is_none()) {
        data.ID = app_id_obj.cast<std::uint32_t>();
    }
    if (!app_key_obj.is_none()) {
        const std::string key = app_key_obj.cast<std::string>();
        std::strncpy(data.encKey, key.c_str(), sizeof(data.encKey) - 1);
    }
    if (!api_level_obj.is_none()) {
        data.version = api_level_obj.cast<std::uint32_t>();
    }

    osdk::ACK::ErrorCode ack = vehicle->activate(&data, kDefaultWaitTimeoutMs);
    return !osdk::ACK::getError(ack);
}

py::dict get_broadcast_real() {
    auto* vehicle = require_vehicle();
    py::dict snapshot;

    auto* broadcast = vehicle->broadcast;
    if (!broadcast) {
        snapshot["attitude"] = make_attitude(1.0, 0.0, 0.0, 0.0);
        snapshot["position"] = make_position(0.0, 0.0, 0.0);
        snapshot["velocity"] = make_velocity(0.0, 0.0, 0.0);
        snapshot["rc"] = make_rc(0, 0, 0, 0);
        snapshot["battery"] = 0.0;
        snapshot["flight_status"] = 0;
        snapshot["display_mode"] = 0;
        snapshot["ctrl_device"] = py::none();
        return snapshot;
    }

    try {
        const auto q = broadcast->getQuaternion();
        snapshot["attitude"] = make_attitude(q.q0, q.q1, q.q2, q.q3);
    } catch (...) {
        snapshot["attitude"] = make_attitude(1.0, 0.0, 0.0, 0.0);
    }

    try {
        const auto pos = broadcast->getGlobalPosition();
        snapshot["position"] = make_position(pos.latitude, pos.longitude, pos.altitude);
    } catch (...) {
        snapshot["position"] = make_position(0.0, 0.0, 0.0);
    }

    try {
        const auto vel = broadcast->getVelocity();
        snapshot["velocity"] = make_velocity(vel.x, vel.y, vel.z);
    } catch (...) {
        snapshot["velocity"] = make_velocity(0.0, 0.0, 0.0);
    }

    try {
        const auto rc = broadcast->getRC();
        snapshot["rc"] = make_rc(rc.pitch, rc.roll, rc.yaw, rc.throttle);
    } catch (...) {
        snapshot["rc"] = make_rc(0, 0, 0, 0);
    }

    try {
        const auto battery = broadcast->getBatteryInfo();
        snapshot["battery"] = battery.percentage;
    } catch (...) {
        snapshot["battery"] = 0.0;
    }

    try {
        const auto status = broadcast->getStatus();
        snapshot["flight_status"] = status.flight;
        snapshot["display_mode"] = status.mode;
    } catch (...) {
        snapshot["flight_status"] = 0;
        snapshot["display_mode"] = 0;
    }

    try {
        const auto sdk = broadcast->getSDKInfo();
        py::dict sdk_info;
        sdk_info["control_mode"] = sdk.controlMode;
        sdk_info["device_status"] = sdk.deviceStatus;
        sdk_info["flight_status"] = sdk.flightStatus;
        snapshot["ctrl_device"] = sdk_info;
    } catch (...) {
        snapshot["ctrl_device"] = py::none();
    }

    return snapshot;
}

bool takeoff_real() {
    auto* vehicle = require_vehicle();
    osdk::ACK::ErrorCode ack = vehicle->control->takeoff(kDefaultWaitTimeoutMs);
    return !osdk::ACK::getError(ack);
}

bool land_real() {
    auto* vehicle = require_vehicle();
    osdk::ACK::ErrorCode ack = vehicle->control->land(kDefaultWaitTimeoutMs);
    return !osdk::ACK::getError(ack);
}

bool rth_real() {
    auto* vehicle = require_vehicle();
    osdk::ACK::ErrorCode ack = vehicle->control->goHome(kDefaultWaitTimeoutMs);
    return !osdk::ACK::getError(ack);
}

bool set_velocity_real(double vx, double vy, double vz, py::object frame_obj, double yaw_rate) {
    auto* vehicle = require_vehicle();
    (void)frame_obj;  // Frame switching handled at a higher level.
    try {
        vehicle->control->velocityAndYawRateCtrl(
            static_cast<float>(vx),
            static_cast<float>(vy),
            static_cast<float>(vz),
            static_cast<float>(yaw_rate));
        return true;
    } catch (...) {
        return false;
    }
}

bool set_attitude_real(double roll, double pitch, double yaw, double thrust, py::object ctrl_flag_obj) {
    auto* vehicle = require_vehicle();
    std::uint8_t flag = static_cast<std::uint8_t>(
        osdk::Control::VERTICAL_THRUST |
        osdk::Control::HORIZONTAL_ANGLE |
        osdk::Control::YAW_ANGLE |
        osdk::Control::HORIZONTAL_BODY);
    if (!ctrl_flag_obj.is_none()) {
        flag = static_cast<std::uint8_t>(ctrl_flag_obj.cast<std::uint32_t>());
    }

    try {
        osdk::Control::CtrlData data(
            flag,
            static_cast<float>(roll),
            static_cast<float>(pitch),
            static_cast<float>(thrust),
            static_cast<float>(yaw));
        vehicle->control->flightCtrl(data);
        return true;
    } catch (...) {
        return false;
    }
}

bool control_management_real(bool obtain) {
    auto* vehicle = require_vehicle();
    osdk::ACK::ErrorCode ack;
    if (obtain) {
        ack = vehicle->obtainCtrlAuthority(kDefaultWaitTimeoutMs);
    } else {
        ack = vehicle->releaseCtrlAuthority(kDefaultWaitTimeoutMs);
    }
    return !osdk::ACK::getError(ack);
}

#else  // DJI_SDK_AVAILABLE

bool g_stub_connected = false;

bool connect_stub(const std::string& device, std::uint32_t baudrate, bool enable_advanced_sensing, bool enable_broadcast) {
    (void)device;
    (void)baudrate;
    (void)enable_advanced_sensing;
    (void)enable_broadcast;
    g_stub_connected = true;
    return g_stub_connected;
}

void disconnect_stub() {
    g_stub_connected = false;
}

bool is_connected_stub() {
    return g_stub_connected;
}

bool activate_stub(py::object app_id, py::object app_key, py::object api_level) {
    (void)app_id;
    (void)app_key;
    (void)api_level;
    return true;
}

py::dict get_broadcast_stub() {
    py::dict d;
    d["attitude"] = make_attitude(1.0, 0.0, 0.0, 0.0);
    d["position"] = make_position(0.0, 0.0, 0.0);
    d["velocity"] = make_velocity(0.0, 0.0, 0.0);
    d["rc"] = make_rc(0, 0, 0, 0);
    d["battery"] = 100.0;
    d["flight_status"] = 0;
    d["display_mode"] = 0;
    d["ctrl_device"] = py::none();
    return d;
}

bool takeoff_stub() { return true; }
bool land_stub() { return true; }
bool rth_stub() { return true; }
bool set_velocity_stub(double vx, double vy, double vz, py::object frame, double yaw_rate) {
    (void)vx;
    (void)vy;
    (void)vz;
    (void)frame;
    (void)yaw_rate;
    return true;
}
bool set_attitude_stub(double roll, double pitch, double yaw, double thrust, py::object ctrl_flag) {
    (void)roll;
    (void)pitch;
    (void)yaw;
    (void)thrust;
    (void)ctrl_flag;
    return true;
}
bool control_management_stub(bool obtain) {
    (void)obtain;
    return true;
}

#endif  // DJI_SDK_AVAILABLE

}  // namespace

PYBIND11_MODULE(_djibindings, m) {
    m.doc() = "DJI SDK binding (conditional real-SDK integration)";
#ifdef DJI_SDK_AVAILABLE
    m.def("connect", &connect_real,
          py::arg("device"), py::arg("baudrate") = 230400,
          py::arg("enable_advanced_sensing") = false,
          py::arg("enable_broadcast") = true,
          "Connect to a DJI vehicle using the Onboard SDK");
    m.def("disconnect", &disconnect_real, "Disconnect and clean up the DJI vehicle instance");
    m.def("is_connected", &is_connected_real, "True if a DJI vehicle instance is active");
    m.def("activate", &activate_real,
          "Activate the DJI vehicle",
          py::arg("app_id") = py::none(),
          py::arg("app_key") = py::none(),
          py::arg("api_level") = py::none());
    m.def("get_broadcast", &get_broadcast_real, "Fetch the latest broadcast telemetry as a dict");
    m.def("takeoff", &takeoff_real, "Initiate a takeoff task");
    m.def("land", &land_real, "Initiate a landing task");
    m.def("return_to_home", &rth_real, "Trigger return-to-home task");
    m.def("set_velocity", &set_velocity_real,
          py::arg("vx"), py::arg("vy"), py::arg("vz"), py::arg("frame") = py::str("local"), py::arg("yaw_rate") = 0.0,
          "Command velocity in body frame (m/s) with yaw rate (rad/s)");
    m.def("set_attitude", &set_attitude_real,
          py::arg("roll"), py::arg("pitch"), py::arg("yaw"), py::arg("thrust"),
          py::arg("ctrl_flag") = py::none(),
          "Command attitude control (Euler angles + thrust)");
    m.def("control_management", &control_management_real,
          py::arg("obtain") = true,
          "Obtain or release control authority");
#else
    m.def("connect", &connect_stub,
          py::arg("device"), py::arg("baudrate") = 230400,
          py::arg("enable_advanced_sensing") = false,
          py::arg("enable_broadcast") = true,
          "Stub connect when DJI SDK is unavailable");
    m.def("disconnect", &disconnect_stub,
          "Stub disconnect when DJI SDK is unavailable");
    m.def("is_connected", &is_connected_stub,
          "Stub is_connected when DJI SDK is unavailable");
    m.def("activate", &activate_stub,
          "Stub activate when DJI SDK is unavailable",
          py::arg("app_id") = py::none(),
          py::arg("app_key") = py::none(),
          py::arg("api_level") = py::none());
    m.def("get_broadcast", &get_broadcast_stub,
          "Stub telemetry when DJI SDK is unavailable");
    m.def("takeoff", &takeoff_stub, "Stub takeoff when DJI SDK is unavailable");
    m.def("land", &land_stub, "Stub land when DJI SDK is unavailable");
    m.def("return_to_home", &rth_stub, "Stub return-to-home when DJI SDK is unavailable");
    m.def("set_velocity", &set_velocity_stub,
          py::arg("vx"), py::arg("vy"), py::arg("vz"), py::arg("frame") = py::str("local"), py::arg("yaw_rate") = 0.0,
          "Stub velocity command when DJI SDK is unavailable");
    m.def("set_attitude", &set_attitude_stub,
          py::arg("roll"), py::arg("pitch"), py::arg("yaw"), py::arg("thrust"),
          py::arg("ctrl_flag") = py::none(),
          "Stub attitude command when DJI SDK is unavailable");
    m.def("control_management", &control_management_stub,
          py::arg("obtain") = true,
          "Stub control management when DJI SDK is unavailable");
#endif
}

