// SPDX-License-Identifier: GPL-2.0-or-later
// Copyright (c) 2026 OokamiKunTV

#include <obs-module.h>
#include <obs-frontend-api.h>
#include <nlohmann/json.hpp>
#include <windows.h>
#include <bcrypt.h>

#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <cmath>
#include <condition_variable>
#include <cstdint>
#include <cstring>
#include <mutex>
#include <optional>
#include <set>
#include <string>
#include <thread>
#include <vector>

#include "bounded_spsc_queue.hpp"

#pragma comment(lib, "bcrypt.lib")

#ifndef LUMI_BRIDGE_VERSION
#define LUMI_BRIDGE_VERSION "0.1.0-development"
#endif

OBS_DECLARE_MODULE()
OBS_MODULE_USE_DEFAULT_LOCALE("lumi-obs-bridge", "en-US")

namespace {
using json = nlohmann::json;
constexpr uint8_t protocol_version = 1;
constexpr size_t max_json_bytes = 64 * 1024;
constexpr size_t max_input_frames = 4096;
constexpr uint32_t output_sample_rate = 16000;
constexpr char bridge_version[] = LUMI_BRIDGE_VERSION;

struct raw_audio_packet {
    std::array<float, max_input_frames> mono{};
    std::array<uint8_t, 16> source_uuid{};
    uint64_t timestamp_ns = 0;
    uint32_t frames = 0;
    uint32_t sample_rate = 48000;
    bool active = false;
    bool muted = false;
};

struct source_description {
    std::string uuid;
    std::string name;
    bool program_active = false;
};

static lumi::bounded_spsc_queue<raw_audio_packet, 32> audio_queue;
static std::condition_variable worker_signal;
static std::mutex worker_signal_mutex;
static std::thread worker_thread;
static std::atomic_bool stopping{false};
static std::atomic_bool source_list_dirty{true};
static std::atomic_bool obs_state_dirty{true};
static std::atomic_uint32_t audio_sequence{0};
static obs_source_t *captured_source = nullptr;
static std::mutex captured_source_mutex;

static std::string wide_to_utf8(const std::wstring &value)
{
    if (value.empty()) return {};
    const int bytes = WideCharToMultiByte(CP_UTF8, 0, value.data(), static_cast<int>(value.size()), nullptr, 0, nullptr, nullptr);
    std::string output(static_cast<size_t>(bytes), '\0');
    WideCharToMultiByte(CP_UTF8, 0, value.data(), static_cast<int>(value.size()), output.data(), bytes, nullptr, nullptr);
    return output;
}

static std::wstring environment_value(const wchar_t *name)
{
    const DWORD needed = GetEnvironmentVariableW(name, nullptr, 0);
    if (!needed) return {};
    std::wstring value(needed - 1, L'\0');
    GetEnvironmentVariableW(name, value.data(), needed);
    return value;
}

static std::string current_user_hash()
{
    const auto identity = wide_to_utf8(environment_value(L"USERDOMAIN") + L"\\" + environment_value(L"USERNAME"));
    BCRYPT_ALG_HANDLE algorithm = nullptr;
    BCRYPT_HASH_HANDLE hash = nullptr;
    std::array<uint8_t, 32> digest{};
    if (BCryptOpenAlgorithmProvider(&algorithm, BCRYPT_SHA256_ALGORITHM, nullptr, 0) < 0) return {};
    DWORD object_size = 0, result_size = 0;
    BCryptGetProperty(algorithm, BCRYPT_OBJECT_LENGTH, reinterpret_cast<PUCHAR>(&object_size), sizeof(object_size), &result_size, 0);
    std::vector<uint8_t> object(object_size);
    const bool ok = BCryptCreateHash(algorithm, &hash, object.data(), object_size, nullptr, 0, 0) >= 0 &&
                    BCryptHashData(hash, reinterpret_cast<PUCHAR>(const_cast<char *>(identity.data())), static_cast<ULONG>(identity.size()), 0) >= 0 &&
                    BCryptFinishHash(hash, digest.data(), static_cast<ULONG>(digest.size()), 0) >= 0;
    if (hash) BCryptDestroyHash(hash);
    BCryptCloseAlgorithmProvider(algorithm, 0);
    if (!ok) return {};
    constexpr char hex[] = "0123456789ABCDEF";
    std::string output;
    output.reserve(16);
    for (size_t index = 0; index < 8; ++index) {
        output.push_back(hex[digest[index] >> 4]);
        output.push_back(hex[digest[index] & 0x0f]);
    }
    return output;
}

static bool parse_uuid(const std::string &value, std::array<uint8_t, 16> &output)
{
    std::string hex;
    for (const auto character : value) if (character != '-') hex.push_back(character);
    if (hex.size() != 32) return false;
    const auto nibble = [](char value) -> int {
        if (value >= '0' && value <= '9') return value - '0';
        if (value >= 'a' && value <= 'f') return value - 'a' + 10;
        if (value >= 'A' && value <= 'F') return value - 'A' + 10;
        return -1;
    };
    for (size_t index = 0; index < output.size(); ++index) {
        const auto high = nibble(hex[index * 2]);
        const auto low = nibble(hex[index * 2 + 1]);
        if (high < 0 || low < 0) return false;
        output[index] = static_cast<uint8_t>((high << 4) | low);
    }
    return true;
}

static void write_u16(std::vector<uint8_t> &value, size_t offset, uint16_t input)
{
    value[offset] = static_cast<uint8_t>(input);
    value[offset + 1] = static_cast<uint8_t>(input >> 8);
}

static void write_u32(std::vector<uint8_t> &value, size_t offset, uint32_t input)
{
    for (size_t index = 0; index < 4; ++index) value[offset + index] = static_cast<uint8_t>(input >> (index * 8));
}

static void write_u64(std::vector<uint8_t> &value, size_t offset, uint64_t input)
{
    for (size_t index = 0; index < 8; ++index) value[offset + index] = static_cast<uint8_t>(input >> (index * 8));
}

static bool write_exact(HANDLE pipe, const uint8_t *data, size_t size)
{
    while (size > 0) {
        DWORD written = 0;
        const auto chunk = static_cast<DWORD>(std::min<size_t>(size, 64 * 1024));
        if (!WriteFile(pipe, data, chunk, &written, nullptr) || written == 0) return false;
        data += written;
        size -= written;
    }
    return true;
}

static bool read_exact(HANDLE pipe, uint8_t *data, size_t size)
{
    while (size > 0) {
        DWORD read = 0;
        if (!ReadFile(pipe, data, static_cast<DWORD>(size), &read, nullptr) || read == 0) return false;
        data += read;
        size -= read;
    }
    return true;
}

static bool write_frame(HANDLE pipe, const uint8_t *data, size_t size)
{
    std::array<uint8_t, 4> prefix{
        static_cast<uint8_t>(size), static_cast<uint8_t>(size >> 8),
        static_cast<uint8_t>(size >> 16), static_cast<uint8_t>(size >> 24)};
    return write_exact(pipe, prefix.data(), prefix.size()) && write_exact(pipe, data, size);
}

static bool write_json(HANDLE pipe, const json &message)
{
    const auto body = message.dump();
    return body.size() <= max_json_bytes && write_frame(pipe, reinterpret_cast<const uint8_t *>(body.data()), body.size());
}

static bool scene_item_active(obs_scene_t *, obs_sceneitem_t *item, void *data);

static void collect_scene_sources(obs_source_t *scene_source, std::set<std::string> &active, std::set<std::string> &visited)
{
    if (!scene_source) return;
    const char *uuid = obs_source_get_uuid(scene_source);
    if (uuid && !visited.insert(uuid).second) return;
    obs_scene_t *scene = obs_scene_from_source(scene_source);
    if (!scene) scene = obs_group_from_source(scene_source);
    if (scene) {
        std::pair values{&active, &visited};
        obs_scene_enum_items(scene, scene_item_active, &values);
    }
}

static bool scene_item_active(obs_scene_t *, obs_sceneitem_t *item, void *data)
{
    if (!obs_sceneitem_visible(item)) return true;
    auto &sets = *static_cast<std::pair<std::set<std::string> *, std::set<std::string> *> *>(data);
    obs_source_t *source = obs_sceneitem_get_source(item);
    if (!source) return true;
    if (const char *uuid = obs_source_get_uuid(source)) sets.first->insert(uuid);
    collect_scene_sources(source, *sets.first, *sets.second);
    return true;
}

static std::set<std::string> program_sources()
{
    std::set<std::string> active, visited;
    obs_source_t *scene = obs_frontend_get_current_scene();
    collect_scene_sources(scene, active, visited);
    if (scene) obs_source_release(scene);
    return active;
}

static bool enumerate_source(void *data, obs_source_t *source)
{
    if (!source || !(obs_source_get_output_flags(source) & OBS_SOURCE_AUDIO)) return true;
    auto &values = *static_cast<std::pair<std::vector<source_description> *, const std::set<std::string> *> *>(data);
    const char *uuid = obs_source_get_uuid(source);
    if (!uuid || !*uuid) return true;
    values.first->push_back({uuid, obs_source_get_name(source) ? obs_source_get_name(source) : "OBS source", values.second->contains(uuid)});
    return true;
}

static std::vector<source_description> enumerate_sources()
{
    const auto active = program_sources();
    std::vector<source_description> sources;
    std::pair values{&sources, &active};
    obs_enum_sources(enumerate_source, &values);
    std::sort(sources.begin(), sources.end(), [](const auto &left, const auto &right) { return left.name < right.name; });
    return sources;
}

static void on_audio(void *, obs_source_t *source, const audio_data *audio, bool muted)
{
    if (!audio || !source || audio->frames == 0 || audio->frames > max_input_frames) return;
    obs_audio_info info{};
    if (!obs_get_audio_info(&info) || !info.samples_per_sec) return;
    const auto channels = std::clamp<uint32_t>(get_audio_channels(info.speakers), 1, MAX_AUDIO_CHANNELS);
    raw_audio_packet packet;
    packet.frames = audio->frames;
    packet.sample_rate = info.samples_per_sec;
    packet.timestamp_ns = audio->timestamp;
    packet.active = obs_source_active(source);
    packet.muted = muted || obs_source_muted(source);
    const char *uuid = obs_source_get_uuid(source);
    if (!uuid || !parse_uuid(uuid, packet.source_uuid)) return;
    for (uint32_t frame = 0; frame < audio->frames; ++frame) {
        float sample = 0.0f;
        for (uint32_t channel = 0; channel < channels; ++channel) {
            if (audio->data[channel]) sample += reinterpret_cast<const float *>(audio->data[channel])[frame];
        }
        packet.mono[frame] = sample / static_cast<float>(channels);
    }
    if (audio_queue.try_push(std::move(packet))) worker_signal.notify_one();
}

static void detach_source()
{
    std::scoped_lock lock(captured_source_mutex);
    if (!captured_source) return;
    obs_source_remove_audio_capture_callback(captured_source, on_audio, nullptr);
    obs_source_release(captured_source);
    captured_source = nullptr;
}

static bool select_source(const std::string &uuid)
{
    detach_source();
    if (uuid.empty()) return true;
    obs_source_t *source = obs_get_source_by_uuid(uuid.c_str());
    if (!source || !(obs_source_get_output_flags(source) & OBS_SOURCE_AUDIO)) {
        if (source) obs_source_release(source);
        return false;
    }
    {
        std::scoped_lock lock(captured_source_mutex);
        captured_source = source;
        obs_source_add_audio_capture_callback(captured_source, on_audio, nullptr);
    }
    return true;
}

static std::vector<uint8_t> encode_audio(const raw_audio_packet &packet)
{
    const auto output_frames = std::min<uint32_t>(3200, static_cast<uint32_t>((static_cast<uint64_t>(packet.frames) * output_sample_rate) / packet.sample_rate));
    std::vector<uint8_t> frame(64 + static_cast<size_t>(output_frames) * 2, 0);
    std::memcpy(frame.data(), "LACP", 4);
    frame[4] = protocol_version;
    frame[5] = static_cast<uint8_t>((packet.active ? 1 : 0) | (packet.muted ? 2 : 0));
    write_u16(frame, 6, 64);
    write_u32(frame, 8, audio_sequence.fetch_add(1, std::memory_order_relaxed));
    write_u64(frame, 12, packet.timestamp_ns / 1000);
    std::copy(packet.source_uuid.begin(), packet.source_uuid.end(), frame.begin() + 36);
    write_u32(frame, 52, output_sample_rate);
    write_u16(frame, 56, 1);
    write_u16(frame, 58, 16);
    write_u32(frame, 60, output_frames * 2);
    for (uint32_t index = 0; index < output_frames; ++index) {
        const auto source_position = std::min<uint32_t>(packet.frames - 1, static_cast<uint32_t>((static_cast<uint64_t>(index) * packet.sample_rate) / output_sample_rate));
        const auto sample = static_cast<int16_t>(std::lrint(std::clamp(packet.mono[source_position], -1.0f, 1.0f) * 32767.0f));
        write_u16(frame, 64 + static_cast<size_t>(index) * 2, static_cast<uint16_t>(sample));
    }
    return frame;
}

static json source_list_message()
{
    json sources = json::array();
    for (const auto &source : enumerate_sources()) {
        sources.push_back({{"source_uuid", source.uuid}, {"display_name", source.name}, {"program_active", source.program_active}, {"source_missing", false}});
    }
    return {{"type", "source_list"}, {"protocol_version", protocol_version}, {"sources", std::move(sources)}};
}

static json obs_state_message()
{
    return {{"type", "obs_state"}, {"protocol_version", protocol_version}, {"version", obs_get_version_string()},
            {"streaming", obs_frontend_streaming_active()}, {"recording", obs_frontend_recording_active()}};
}

static bool output_caption(const std::string &text, double display_seconds)
{
    if (!obs_frontend_streaming_active() || text.empty()) return false;
    obs_output_t *output = obs_frontend_get_streaming_output();
    if (!output) return false;
    obs_output_output_caption_text2(output, text.c_str(), std::clamp(display_seconds, 0.5, 10.0));
    obs_output_release(output);
    return true;
}

static std::optional<json> handle_command(const json &message)
{
    const auto type = message.value("type", "");
    if (type == "select_sources") {
        const auto uuid = message.value("primary_source_uuid", "");
        const auto installed = select_source(uuid);
        source_list_dirty.store(true, std::memory_order_release);
        blog(installed ? LOG_INFO : LOG_WARNING, "[Lumi Companion] %s selected OBS audio source %s",
             installed ? "Attached" : "Could not attach", uuid.c_str());
        return json{{"type", "selection_state"}, {"protocol_version", protocol_version}, {"source_uuid", uuid}, {"attached", installed}};
    } else if (type == "caption" && message.contains("payload")) {
        const auto &payload = message["payload"];
        const auto text = payload.value("stable_text", "");
        const auto duration = payload.value("display_seconds", 2.0);
        output_caption(text, duration);
    }
    return std::nullopt;
}

static bool read_available_command(HANDLE pipe)
{
    DWORD available = 0;
    if (!PeekNamedPipe(pipe, nullptr, 0, nullptr, &available, nullptr)) return false;
    if (available < 4) return true;
    std::array<uint8_t, 4> prefix{};
    if (!read_exact(pipe, prefix.data(), prefix.size())) return false;
    const uint32_t size = static_cast<uint32_t>(prefix[0]) | (static_cast<uint32_t>(prefix[1]) << 8) |
                          (static_cast<uint32_t>(prefix[2]) << 16) | (static_cast<uint32_t>(prefix[3]) << 24);
    if (size == 0 || size > max_json_bytes) return false;
    std::vector<uint8_t> body(size);
    if (!read_exact(pipe, body.data(), body.size())) return false;
    try {
        const auto response = handle_command(json::parse(body.begin(), body.end()));
        if (response && !write_json(pipe, *response)) return false;
    }
    catch (const std::exception &error) { blog(LOG_WARNING, "[Lumi Companion] Ignored invalid IPC command: %s", error.what()); }
    return true;
}

static HANDLE connect_pipe()
{
    const auto hash = current_user_hash();
    if (hash.empty()) return INVALID_HANDLE_VALUE;
    const auto name = L"\\\\.\\pipe\\Lumi.Companion.ObsBridge.v1." + std::wstring(hash.begin(), hash.end());
    return CreateFileW(name.c_str(), GENERIC_READ | GENERIC_WRITE, 0, nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
}

static void run_pipe_worker()
{
    while (!stopping.load(std::memory_order_acquire)) {
        HANDLE pipe = connect_pipe();
        if (pipe == INVALID_HANDLE_VALUE) {
            std::unique_lock lock(worker_signal_mutex);
            worker_signal.wait_for(lock, std::chrono::seconds(2));
            continue;
        }
        blog(LOG_INFO, "[Lumi Companion] connected to the same-user Companion IPC endpoint");
        bool connected = write_json(pipe, {{"type", "hello"}, {"protocol_version", protocol_version}, {"obs_version", obs_get_version_string()}, {"bridge_version", bridge_version}});
        source_list_dirty.store(true, std::memory_order_release);
        obs_state_dirty.store(true, std::memory_order_release);
        auto health_sent = std::chrono::steady_clock::now();
        while (connected && !stopping.load(std::memory_order_acquire)) {
            connected = read_available_command(pipe);
            if (connected && source_list_dirty.exchange(false, std::memory_order_acq_rel)) connected = write_json(pipe, source_list_message());
            if (connected && obs_state_dirty.exchange(false, std::memory_order_acq_rel)) connected = write_json(pipe, obs_state_message());
            while (connected) {
                auto packet = audio_queue.try_pop();
                if (!packet) break;
                const auto encoded = encode_audio(*packet);
                connected = write_frame(pipe, encoded.data(), encoded.size());
            }
            const auto now = std::chrono::steady_clock::now();
            if (connected && now - health_sent >= std::chrono::seconds(10)) {
                connected = write_json(pipe, {{"type", "health"}, {"protocol_version", protocol_version},
                    {"status", "healthy"}, {"audio_frames_dropped", audio_queue.dropped()}});
                health_sent = now;
            }
            std::unique_lock lock(worker_signal_mutex);
            worker_signal.wait_for(lock, std::chrono::milliseconds(10));
        }
        CloseHandle(pipe);
        blog(LOG_INFO, "[Lumi Companion] disconnected from Companion IPC; retrying safely");
    }
}

static void frontend_event(obs_frontend_event event, void *)
{
    switch (event) {
    case OBS_FRONTEND_EVENT_STREAMING_STARTED:
    case OBS_FRONTEND_EVENT_STREAMING_STOPPED:
    case OBS_FRONTEND_EVENT_RECORDING_STARTED:
    case OBS_FRONTEND_EVENT_RECORDING_STOPPED:
        obs_state_dirty.store(true, std::memory_order_release);
        break;
    case OBS_FRONTEND_EVENT_SCENE_CHANGED:
    case OBS_FRONTEND_EVENT_SCENE_LIST_CHANGED:
    case OBS_FRONTEND_EVENT_SCENE_COLLECTION_CHANGED:
    case OBS_FRONTEND_EVENT_FINISHED_LOADING:
        source_list_dirty.store(true, std::memory_order_release);
        break;
    default:
        break;
    }
    worker_signal.notify_one();
}
} // namespace

const char *obs_module_description(void)
{
    return "Companion-managed Lumi audio and native-caption bridge";
}

bool obs_module_load(void)
{
    const char *version = obs_get_version_string();
    const int major = version ? std::atoi(version) : 0;
    if (major < 31) {
        blog(LOG_ERROR, "[Lumi Companion] OBS %s is unsupported; version 31 or newer is required", version ? version : "unknown");
        return false;
    }
    stopping.store(false, std::memory_order_release);
    obs_frontend_add_event_callback(frontend_event, nullptr);
    worker_thread = std::thread(run_pipe_worker);
    blog(LOG_INFO, "[Lumi Companion] bridge %s loaded", bridge_version);
    return true;
}

void obs_module_unload(void)
{
    obs_frontend_remove_event_callback(frontend_event, nullptr);
    detach_source();
    stopping.store(true, std::memory_order_release);
    worker_signal.notify_all();
    if (worker_thread.joinable()) worker_thread.join();
    blog(LOG_INFO, "[Lumi Companion] bridge unloaded");
}