#include <napi.h>

#include <cerrno>
#include <cstring>
#include <string>
#include <vector>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

/* ── IPC protocol constants ──────────────────────────────────────────── */

static constexpr uint32_t IPC_MAGIC   = 0x4D564741;
static constexpr uint16_t IPC_VERSION = 1;

enum MsgType : uint16_t {
    MSG_QUERY_DEVICES       = 0x0001,
    MSG_SUBMIT_WORKLOAD     = 0x0002,
    MSG_ALLOC_MEMORY        = 0x0003,
    MSG_FREE_MEMORY         = 0x0004,
    MSG_QUERY_CAPABILITIES  = 0x0005,
    MSG_SET_STRATEGY        = 0x0006,
    MSG_SUBSCRIBE_TELEMETRY = 0x0007,
};

/* Wire header: 16 bytes total */
#pragma pack(push, 1)
struct IpcHeader {
    uint32_t magic;
    uint16_t version;
    uint16_t msg_type;
    uint32_t payload_len;
    uint32_t seq;
};
#pragma pack(pop)

/* ── MvgalClient ─────────────────────────────────────────────────────── */

class MvgalClient : public Napi::ObjectWrap<MvgalClient> {
public:
    static Napi::Object Init(Napi::Env env, Napi::Object exports);
    explicit MvgalClient(const Napi::CallbackInfo &info);

private:
    ~MvgalClient();

    /* JS-callable methods */
    Napi::Value Close(const Napi::CallbackInfo &info);
    Napi::Value QueryDevices(const Napi::CallbackInfo &info);
    Napi::Value QueryCapabilities(const Napi::CallbackInfo &info);
    Napi::Value SubmitWorkload(const Napi::CallbackInfo &info);
    Napi::Value SetStrategy(const Napi::CallbackInfo &info);
    Napi::Value AllocateMemory(const Napi::CallbackInfo &info);
    Napi::Value FreeMemory(const Napi::CallbackInfo &info);
    Napi::Value SubscribeTelemetry(const Napi::CallbackInfo &info);
    Napi::Value HealthCheck(const Napi::CallbackInfo &info);

    /* Internal helpers */
    Napi::Value SendMessage(Napi::Env env, MsgType type,
                             const std::string &payload_json);
    std::string RecvExact(size_t size);
    void EnsureConnected();

    int  sock_fd_ = -1;
    int  seq_     = 0;
    bool closed_  = false;
};

/* ── Constructor / Destructor ────────────────────────────────────────── */

MvgalClient::MvgalClient(const Napi::CallbackInfo &info)
    : Napi::ObjectWrap<MvgalClient>(info)
{
    Napi::Env env = info.Env();

    std::string socket_path = "/run/mvgal/mvgal.sock";
    if (info.Length() > 0 && info[0].IsString()) {
        socket_path = info[0].As<Napi::String>().Utf8Value();
    }

    sock_fd_ = ::socket(AF_UNIX, SOCK_STREAM, 0);
    if (sock_fd_ < 0) {
        Napi::Error::New(env, "Failed to create socket: " +
                              std::string(std::strerror(errno)))
            .ThrowAsJavaScriptException();
        return;
    }

    struct sockaddr_un addr;
    std::memset(&addr, 0, sizeof(addr));
    addr.sun_family = AF_UNIX;
    std::strncpy(addr.sun_path, socket_path.c_str(), sizeof(addr.sun_path) - 1);

    if (::connect(sock_fd_, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        ::close(sock_fd_);
        sock_fd_ = -1;
        Napi::Error::New(env, "Cannot connect to mvgal at " + socket_path +
                              ": " + std::string(std::strerror(errno)))
            .ThrowAsJavaScriptException();
        return;
    }
}

MvgalClient::~MvgalClient() {
    if (sock_fd_ >= 0) {
        ::close(sock_fd_);
        sock_fd_ = -1;
    }
}

/* ── Internal helpers ────────────────────────────────────────────────── */

void MvgalClient::EnsureConnected() {
    if (closed_ || sock_fd_ < 0) {
        throw std::runtime_error("Not connected to mvgal daemon");
    }
}

std::string MvgalClient::RecvExact(size_t size) {
    std::string buf(size, '\0');
    size_t offset = 0;
    while (offset < size) {
        ssize_t n = ::read(sock_fd_, &buf[offset], size - offset);
        if (n <= 0) {
            closed_ = true;
            throw std::runtime_error("Connection closed by peer");
        }
        offset += static_cast<size_t>(n);
    }
    return buf;
}

Napi::Value MvgalClient::SendMessage(Napi::Env env, MsgType type,
                                      const std::string &payload_json)
{
    EnsureConnected();

    uint32_t payload_len = static_cast<uint32_t>(payload_json.size());
    seq_++;

    IpcHeader hdr;
    hdr.magic       = IPC_MAGIC;
    hdr.version     = IPC_VERSION;
    hdr.msg_type    = static_cast<uint16_t>(type);
    hdr.payload_len = payload_len;
    hdr.seq         = static_cast<uint32_t>(seq_);

    /* Send header + payload */
    {
        std::vector<struct iovec> iov(2);
        iov[0].iov_base = &hdr;
        iov[0].iov_len  = sizeof(hdr);
        iov[1].iov_base = const_cast<char *>(payload_json.data());
        iov[1].iov_len  = payload_len;

        struct msghdr msg;
        std::memset(&msg, 0, sizeof(msg));
        msg.msg_iov    = iov.data();
        msg.msg_iovlen = 2;

        if (::sendmsg(sock_fd_, &msg, 0) < 0) {
            closed_ = true;
            Napi::Error::New(env, "Send failed: " +
                                  std::string(std::strerror(errno)))
                .ThrowAsJavaScriptException();
            return env.Undefined();
        }
    }

    /* Read response header */
    IpcHeader resp;
    try {
        std::string raw = RecvExact(sizeof(resp));
        std::memcpy(&resp, raw.data(), sizeof(resp));
    } catch (const std::exception &e) {
        Napi::Error::New(env, e.what()).ThrowAsJavaScriptException();
        return env.Undefined();
    }

    if (resp.magic != IPC_MAGIC) {
        Napi::Error::New(env, "Invalid response magic").ThrowAsJavaScriptException();
        return env.Undefined();
    }

    /* Read response payload */
    if (resp.payload_len > 0) {
        try {
            std::string body = RecvExact(resp.payload_len);
            return Napi::JSON::Parse(env, body);
        } catch (const std::exception &e) {
            Napi::Error::New(env, e.what()).ThrowAsJavaScriptException();
            return env.Undefined();
        }
    }

    return env.Undefined();
}

/* ── JS-callable methods ─────────────────────────────────────────────── */

Napi::Value MvgalClient::Close(const Napi::CallbackInfo &info) {
    if (sock_fd_ >= 0) {
        ::close(sock_fd_);
        sock_fd_ = -1;
    }
    closed_ = true;
    return info.Env().Undefined();
}

Napi::Value MvgalClient::QueryDevices(const Napi::CallbackInfo &info) {
    return SendMessage(info.Env(), MSG_QUERY_DEVICES, "");
}

Napi::Value MvgalClient::QueryCapabilities(const Napi::CallbackInfo &info) {
    return SendMessage(info.Env(), MSG_QUERY_CAPABILITIES, "");
}

Napi::Value MvgalClient::SubmitWorkload(const Napi::CallbackInfo &info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsObject()) {
        Napi::TypeError::New(env, "Expected an object argument")
            .ThrowAsJavaScriptException();
        return env.Undefined();
    }
    std::string json = Napi::JSON::Stringify(env, info[0].As<Napi::Object>());
    return SendMessage(env, MSG_SUBMIT_WORKLOAD, json);
}

Napi::Value MvgalClient::SetStrategy(const Napi::CallbackInfo &info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsString()) {
        Napi::TypeError::New(env, "Expected a strategy string argument")
            .ThrowAsJavaScriptException();
        return env.Undefined();
    }
    std::string strategy = info[0].As<Napi::String>().Utf8Value();
    std::string json = "{\"strategy\":\"" + strategy + "\"}";
    return SendMessage(env, MSG_SET_STRATEGY, json);
}

Napi::Value MvgalClient::AllocateMemory(const Napi::CallbackInfo &info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsNumber()) {
        Napi::TypeError::New(env, "Expected a size (number) argument")
            .ThrowAsJavaScriptException();
        return env.Undefined();
    }
    uint64_t size = info[0].As<Napi::Number>().Uint64Value();
    std::string json = "{\"size\":" + std::to_string(size) + "}";
    return SendMessage(env, MSG_ALLOC_MEMORY, json);
}

Napi::Value MvgalClient::FreeMemory(const Napi::CallbackInfo &info) {
    Napi::Env env = info.Env();
    if (info.Length() < 1 || !info[0].IsString()) {
        Napi::TypeError::New(env, "Expected an allocation_id string argument")
            .ThrowAsJavaScriptException();
        return env.Undefined();
    }
    std::string id = info[0].As<Napi::String>().Utf8Value();
    std::string json = "{\"allocation_id\":\"" + id + "\"}";
    return SendMessage(env, MSG_FREE_MEMORY, json);
}

Napi::Value MvgalClient::SubscribeTelemetry(const Napi::CallbackInfo &info) {
    Napi::Env env = info.Env();
    uint64_t interval_ms = 1000;
    if (info.Length() > 0 && info[0].IsNumber()) {
        interval_ms = info[0].As<Napi::Number>().Uint64Value();
    }
    std::string json = "{\"interval_ms\":" + std::to_string(interval_ms) + "}";
    return SendMessage(env, MSG_SUBSCRIBE_TELEMETRY, json);
}

Napi::Value MvgalClient::HealthCheck(const Napi::CallbackInfo &info) {
    Napi::Env env = info.Env();
    try {
        EnsureConnected();
        SendMessage(env, MSG_QUERY_CAPABILITIES, "");
        return Napi::Boolean::New(env, true);
    } catch (...) {
        return Napi::Boolean::New(env, false);
    }
}

/* ── Module init ─────────────────────────────────────────────────────── */

Napi::Object MvgalClient::Init(Napi::Env env, Napi::Object exports) {
    Napi::Function func = DefineClass(env, "MvgalClient", {
        InstanceMethod("close",              &MvgalClient::Close),
        InstanceMethod("queryDevices",       &MvgalClient::QueryDevices),
        InstanceMethod("queryCapabilities",  &MvgalClient::QueryCapabilities),
        InstanceMethod("submitWorkload",     &MvgalClient::SubmitWorkload),
        InstanceMethod("setStrategy",        &MvgalClient::SetStrategy),
        InstanceMethod("allocateMemory",     &MvgalClient::AllocateMemory),
        InstanceMethod("freeMemory",         &MvgalClient::FreeMemory),
        InstanceMethod("subscribeTelemetry", &MvgalClient::SubscribeTelemetry),
        InstanceMethod("healthCheck",        &MvgalClient::HealthCheck),
    });

    exports.Set("MvgalClient", func);
    return exports;
}

NODE_API_MODULE(mvgal, MvgalClient::Init)
