accelerator/include/ipc_handler.hpp

//
// SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
// SPDX-License-Identifier: Apache-2.0
//

#ifndef POCKET_TTS_IPC_HANDLER_HPP
#define POCKET_TTS_IPC_HANDLER_HPP

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <vector>

namespace pocket_tts_accelerator {

enum class CommandType : std::uint32_t {
    PING = 0,
    PROCESS_AUDIO = 1,
    CONVERT_TO_MONO = 2,
    CONVERT_TO_PCM = 3,
    RESAMPLE_AUDIO = 4,
    GET_MEMORY_STATS = 5,
    CLEAR_MEMORY_POOL = 6,
    SHUTDOWN = 7,
    UNKNOWN = 255
};

enum class ResponseStatus : std::uint32_t {
    SUCCESS = 0,
    ERROR_INVALID_COMMAND = 1,
    ERROR_FILE_NOT_FOUND = 2,
    ERROR_PROCESSING_FAILED = 3,
    ERROR_MEMORY_ALLOCATION = 4,
    ERROR_INTERNAL = 5
};

struct RequestHeader {
    std::uint32_t magic_number;
    std::uint32_t command_type;
    std::uint32_t payload_size;
    std::uint32_t request_id;
};

struct ResponseHeader {
    std::uint32_t magic_number;
    std::uint32_t status_code;
    std::uint32_t payload_size;
    std::uint32_t request_id;
};

struct ProcessAudioRequest {
    char input_file_path[512];
    char output_file_path[512];
    std::uint32_t target_sample_rate;
    std::uint32_t options_flags;
};

struct MemoryStatsResponse {
    std::uint64_t total_allocated_bytes;
    std::uint64_t total_used_bytes;
    std::uint64_t block_count;
};

class IpcHandler {
public:
    using CommandHandlerFunction = std::function<std::vector<std::uint8_t>(const std::vector<std::uint8_t>&)>;

    explicit IpcHandler(const std::string& socket_path);
    ~IpcHandler();

    IpcHandler(const IpcHandler&) = delete;
    IpcHandler& operator=(const IpcHandler&) = delete;

    bool start_server();
    void stop_server();
    bool is_running() const;

    void register_command_handler(CommandType command_type, CommandHandlerFunction handler);
    void set_shutdown_callback(std::function<void()> callback);

    static constexpr std::uint32_t PROTOCOL_MAGIC_NUMBER = 0x50545453;
    static constexpr std::size_t MAXIMUM_PAYLOAD_SIZE = 16 * 1024 * 1024;
    static constexpr int CONNECTION_BACKLOG = 5;

private:
    void accept_connections_loop();
    void handle_client_connection(int client_socket_fd);
    bool send_response(int socket_fd, const ResponseHeader& header, const std::vector<std::uint8_t>& payload);
    bool receive_request(int socket_fd, RequestHeader& header, std::vector<std::uint8_t>& payload);

    std::string socket_file_path;
    int server_socket_fd;
    std::atomic<bool> is_server_running;
    std::thread accept_thread;
    std::mutex handlers_mutex;
    std::unordered_map<CommandType, CommandHandlerFunction> command_handlers;
    std::function<void()> shutdown_callback;
};

}

#endif