/**
 * @file genai_demo.cpp
 * @brief Standalone demonstration of GenAI module architecture
 *
 * @par Architecture Overview
 *
 * This program demonstrates a thread-pool based GenAI module designed for
 * integration into ProxySQL. The architecture follows these principles:
 *
 * - **Zero Coupling**: GenAI only knows about file descriptors, not client types
 * - **Non-Blocking**: Clients never wait on GenAI responses
 * - **Scalable**: Fixed resource usage (bounded thread pool)
 * - **Observable**: Easy to monitor and debug
 * - **Testable**: Standalone, independent testing
 *
 * @par Components
 *
 * 1. **GenAI Module** (GenAIModule class)
 *    - Listener thread using epoll to monitor all client file descriptors
 *    - Thread-safe request queue with condition variable
 *    - Fixed-size worker thread pool (configurable, default 4)
 *    - Processes requests asynchronously without blocking clients
 *
 * 2. **Client** (Client class)
 *    - Simulates MySQL/PgSQL threads in ProxySQL
 *    - Creates socketpair connections to GenAI
 *    - Sends requests non-blocking
 *    - Polls for responses asynchronously
 *
 * @par Communication Flow
 *
 * @msc
 * Client, GenAI_Listener, GenAI_Worker;
 *
 * Client note Client
 * Client->GenAI_Listener: socketpair() creates 2 FDs;
 * Client->GenAI_Listener: register_client(write_fd);
 * Client->GenAI_Listener: send request (async);
 * Client note Client continues working;
 * GenAI_Listener>>GenAI_Worker: enqueue request;
 * GenAI_Worker note GenAI_Worker processes (simulate with sleep);
 * GenAI_Worker->Client: write response to socket;
 * Client note Client receives response when polling;
 * @endmsc
 *
 * @par Build and Run
 *
 * @code{.sh}
 * # Compile
 * g++ -std=c++17 -o genai_demo genai_demo.cpp -lpthread
 *
 * # Run
 * ./genai_demo
 *
 * # Or use the Makefile
 * make run
 * @endcode
 *
 * @author ProxySQL Team
 * @date 2025-01-08
 * @version 1.0
 */

#include <iostream>
#include <thread>
#include <vector>
#include <queue>
#include <mutex>
#include <condition_variable>
#include <unordered_map>
#include <unordered_set>
#include <atomic>
#include <cstring>
#include <unistd.h>
#include <fcntl.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/eventfd.h>
#include <chrono>
#include <random>

// ============================================================================
// Platform Compatibility
// ============================================================================

/**
 * @def EFD_CLOEXEC
 * @brief Close-on-exec flag for eventfd()
 *
 * Set the close-on-exec (FD_CLOEXEC) flag on the new file descriptor.
 * This ensures the file descriptor is automatically closed when exec() is called.
 */
#ifndef EFD_CLOEXEC
#define EFD_CLOEXEC 0200000
#endif

/**
 * @def EFD_NONBLOCK
 * @brief Non-blocking flag for eventfd()
 *
 * Set the O_NONBLOCK flag on the new file descriptor.
 * This allows read() and write() operations to return immediately with EAGAIN
 * if the operation would block.
 */
#ifndef EFD_NONBLOCK
#define EFD_NONBLOCK 04000
#endif

// ============================================================================
// Protocol Definitions
// ============================================================================

/**
 * @struct RequestHeader
 * @brief Header structure for client requests to GenAI module
 *
 * This structure is sent first, followed by input_size bytes of input data.
 * All fields are in network byte order (little-endian on x86_64).
 *
 * @var RequestHeader::request_id
 * Unique identifier for this request. Generated by the client and echoed
 * back in the response for correlation. Allows tracking of multiple
 * concurrent requests from the same client.
 *
 * @var RequestHeader::operation
 * Type of operation to perform. See Operation enum for valid values:
 * - OP_EMBEDDING (0): Generate text embeddings
 * - OP_COMPLETION (1): Text completion/generation
 * - OP_RAG (2): Retrieval-augmented generation
 *
 * @var RequestHeader::input_size
 * Size in bytes of the input data that follows this header.
 * The input data is sent immediately after this header.
 *
 * @var RequestHeader::flags
 * Reserved for future use. Set to 0.
 *
 * @par Example
 * @code{.cpp}
 * RequestHeader req;
 * req.request_id = 12345;
 * req.operation = OP_EMBEDDING;
 * req.input_size = text.length();
 * req.flags = 0;
 *
 * write(fd, &req, sizeof(req));
 * write(fd, text.data(), text.length());
 * @endcode
 */
struct RequestHeader {
    uint64_t request_id;     ///< Unique request identifier for correlation
    uint32_t operation;      ///< Operation type (OP_EMBEDDING, OP_COMPLETION, OP_RAG)
    uint32_t input_size;     ///< Size of input data following header (bytes)
    uint32_t flags;          ///< Reserved for future use (set to 0)
};

/**
 * @struct ResponseHeader
 * @brief Header structure for GenAI module responses to clients
 *
 * This structure is sent first, followed by output_size bytes of output data.
 *
 * @var ResponseHeader::request_id
 * Echoes the request_id from the original RequestHeader.
 * Used by the client to correlate the response with the pending request.
 *
 * @var ResponseHeader::status_code
 * Status of the request processing:
 * - 0: Success
 * - >0: Error code (specific codes to be defined)
 *
 * @var ResponseHeader::output_size
 * Size in bytes of the output data that follows this header.
 * May be 0 if there is no output data.
 *
 * @var ResponseHeader::processing_time_ms
 * Time taken by GenAI to process this request, in milliseconds.
 * Useful for performance monitoring and debugging.
 *
 * @par Example
 * @code{.cpp}
 * ResponseHeader resp;
 * read(fd, &resp, sizeof(resp));
 *
 * std::vector<char> output(resp.output_size);
 * read(fd, output.data(), resp.output_size);
 *
 * if (resp.status_code == 0) {
 *     std::cout << "Processed in " << resp.processing_time_ms << "ms\n";
 * }
 * @endcode
 */
struct ResponseHeader {
    uint64_t request_id;         ///< Echo of client's request identifier
    uint32_t status_code;        ///< 0=success, >0=error
    uint32_t output_size;        ///< Size of output data following header (bytes)
    uint32_t processing_time_ms; ///< Actual processing time in milliseconds
};

/**
 * @enum Operation
 * @brief Supported GenAI operations
 *
 * Defines the types of operations that the GenAI module can perform.
 * These values are used in the RequestHeader::operation field.
 *
 * @var OP_EMBEDDING
 * Generate text embeddings using an embedding model.
 * Input: Text string to embed
 * Output: Vector of floating-point numbers (the embedding)
 *
 * @var OP_COMPLETION
 * Generate text completion using a language model.
 * Input: Prompt text
 * Output: Generated completion text
 *
 * @var OP_RAG
 * Retrieval-augmented generation.
 * Input: Query text
 * Output: Generated response with retrieved context
 */
enum Operation {
    OP_EMBEDDING = 0,    ///< Generate text embeddings (e.g., OpenAI text-embedding-3-small)
    OP_COMPLETION = 1,   ///< Text completion/generation (e.g., GPT-4)
    OP_RAG = 2           ///< Retrieval-augmented generation
};

// ============================================================================
// GenAI Module
// ============================================================================

/**
 * @class GenAIModule
 * @brief Thread-pool based GenAI processing module
 *
 * The GenAI module implements an asynchronous request processing system
 * designed to handle GenAI operations (embeddings, completions, RAG) without
 * blocking calling threads.
 *
 * @par Architecture
 *
 * - **Listener Thread**: Uses epoll to monitor all client file descriptors.
 *   When data arrives on any FD, it reads the request, validates it, and
 *   pushes it onto the request queue.
 *
 * - **Request Queue**: Thread-safe FIFO queue that holds pending requests.
 *   Protected by a mutex and signaled via condition variable.
 *
 * - **Worker Threads**: Fixed-size thread pool (configurable) that processes
 *   requests from the queue. Each worker waits for work, processes requests
 *   (potentially blocking on I/O to external services), and writes responses
 *   back to clients.
 *
 * @par Threading Model
 *
 * - One listener thread (epoll-based I/O multiplexing)
 * - N worker threads (configurable via constructor)
 * - Total threads = 1 + num_workers
 *
 * @par Thread Safety
 *
 * - Public methods are thread-safe
 * - Multiple clients can register/unregister concurrently
 * - Request queue is protected by mutex
 * - Client FD set is protected by mutex
 *
 * @par Usage Example
 * @code{.cpp}
 * // Create GenAI module with 8 workers
 * GenAIModule genai(8);
 *
 * // Start the module (spawns threads)
 * genai.start();
 *
 * // Register clients
 * int client_fd = get_socket_from_client();
 * genai.register_client(client_fd);
 *
 * // ... module runs, processing requests ...
 *
 * // Shutdown
 * genai.stop();
 * @endcode
 *
 * @par Shutdown Sequence
 *
 * 1. Set running_ flag to false
 * 2. Write to event_fd to wake listener
 * 3. Notify all worker threads via condition variable
 * 4. Join all threads
 * 5. Close all client FDs
 * 6. Close epoll and event FDs
 */
class GenAIModule {
public:

    /**
     * @struct Request
     * @brief Internal request structure queued for worker processing
     *
     * This structure represents a request after it has been received from
     * a client and enqueued for processing by worker threads.
     *
     * @var Request::client_fd
     * File descriptor to write the response to. This is the client's
     * socket FD that was registered via register_client().
     *
     * @var Request::request_id
     * Client's request identifier for correlation. Echoed back in response.
     *
     * @var Request::operation
     * Operation type (OP_EMBEDDING, OP_COMPLETION, or OP_RAG).
     *
     * @var Request::input
     * Input data (text prompt, etc.) from the client.
     */
    struct Request {
        int client_fd;              ///< Where to send response (client's socket FD)
        uint64_t request_id;        ///< Client's correlation identifier
        uint32_t operation;         ///< Type of operation to perform
        std::string input;          ///< Input data (text, prompt, etc.)
    };

    /**
     * @brief Construct a GenAI module with specified worker count
     *
     * Creates a GenAI module instance. The module is not started until
     * start() is called.
     *
     * @param num_workers Number of worker threads in the pool (default: 4)
     *
     * @par Worker Count Guidelines
     * - For I/O-bound operations (API calls): 4-8 workers typically sufficient
     * - For CPU-bound operations (local models): match CPU core count
     * - Too many workers can cause contention; too few can cause queue buildup
     */
    GenAIModule(int num_workers = 4)
        : num_workers_(num_workers), running_(false) {}

    /**
     * @brief Start the GenAI module
     *
     * Initializes and starts all internal threads:
     * - Creates epoll instance for listener
     * - Creates eventfd for shutdown signaling
     * - Spawns worker threads (each runs worker_loop())
     * - Spawns listener thread (runs listener_loop())
     *
     * @post Module is running and ready to accept clients
     * @post All worker threads are waiting for requests
     * @post Listener thread is monitoring registered client FDs
     *
     * @note This method blocks briefly during thread creation but returns
     *       immediately after threads are spawned.
     *
     * @warning Do not call start() on an already-running module.
     *          Call stop() first.
     *
     * @par Thread Creation Sequence
     * 1. Create epoll instance for I/O multiplexing
     * 2. Create eventfd for shutdown notification
     * 3. Add eventfd to epoll set
     * 4. Spawn N worker threads (each calls worker_loop())
     * 5. Spawn listener thread (calls listener_loop())
     */
    void start() {
        running_ = true;

        // Create epoll instance for listener
        // EPOLL_CLOEXEC: Close on exec to prevent FD leaks
        epoll_fd_ = epoll_create1(EPOLL_CLOEXEC);
        if (epoll_fd_ < 0) {
            perror("epoll_create1");
            exit(1);
        }

        // Create eventfd for shutdown notification
        // EFD_NONBLOCK: Non-blocking operations
        // EFD_CLOEXEC: Close on exec
        event_fd_ = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC);
        if (event_fd_ < 0) {
            perror("eventfd");
            exit(1);
        }

        // Add eventfd to epoll set so listener can be woken for shutdown
        struct epoll_event ev;
        ev.events = EPOLLIN;
        ev.data.fd = event_fd_;
        if (epoll_ctl(epoll_fd_, EPOLL_CTL_ADD, event_fd_, &ev) < 0) {
            perror("epoll_ctl eventfd");
            exit(1);
        }

        // Start worker threads
        // Each worker runs worker_loop() and waits for requests
        for (int i = 0; i < num_workers_; i++) {
            worker_threads_.emplace_back([this, i]() { worker_loop(i); });
        }

        // Start listener thread
        // Listener runs listener_loop() and monitors client FDs via epoll
        listener_thread_ = std::thread([this]() { listener_loop(); });

        std::cout << "[GenAI] Module started with " << num_workers_ << " workers\n";
    }

    /**
     * @brief Register a new client connection with the GenAI module
     *
     * Registers a client's file descriptor with the epoll set so the
     * listener thread can monitor it for incoming requests.
     *
     * @param client_fd File descriptor to monitor (one end of socketpair)
     *
     * @pre client_fd is a valid, open file descriptor
     * @pre Module has been started (start() was called)
     * @pre client_fd has not already been registered
     *
     * @post client_fd is added to epoll set for monitoring
     * @post client_fd is set to non-blocking mode
     * @post Listener thread will be notified when data arrives on client_fd
     *
     * @par Thread Safety
     * This method is thread-safe and can be called concurrently by
     * multiple threads registering different clients.
     *
     * @par Client Registration Flow
     * 1. Client creates socketpair() (2 FDs)
     * 2. Client keeps one FD for reading responses
     * 3. Client passes other FD to this method
     * 4. This FD is added to epoll set
     * 5. When client writes request, listener is notified
     *
     * @note This method is typically called by the client after creating
     *       a socketpair(). The client keeps one end, the GenAI module
     *       gets the other end.
     *
     * @warning The caller retains ownership of the original socketpair FDs
     *          and is responsible for closing them after unregistering.
     */
    void register_client(int client_fd) {
        std::lock_guard<std::mutex> lock(clients_mutex_);

        // Set FD to non-blocking mode
        // This ensures read/write operations don't block the listener
        int flags = fcntl(client_fd, F_GETFL, 0);
        fcntl(client_fd, F_SETFL, flags | O_NONBLOCK);

        // Add to epoll set
        // EPOLLIN: Notify when FD is readable (data available to read)
        struct epoll_event ev;
        ev.events = EPOLLIN;
        ev.data.fd = client_fd;
        if (epoll_ctl(epoll_fd_, EPOLL_CTL_ADD, client_fd, &ev) < 0) {
            perror("epoll_ctl client_fd");
            return;
        }

        client_fds_.insert(client_fd);
        std::cout << "[GenAI] Registered client fd " << client_fd << "\n";
    }

    /**
     * @brief Stop the GenAI module and clean up resources
     *
     * Initiates graceful shutdown:
     * - Sets running_ flag to false (signals threads to stop)
     * - Writes to event_fd to wake listener from epoll_wait()
     * - Notifies all workers via condition variable
     * - Joins all threads (waits for them to finish)
     * - Closes all client file descriptors
     * - Closes epoll and event FDs
     *
     * @post All threads have stopped
     * @post All resources are cleaned up
     * @post Module can be restarted by calling start() again
     *
     * @par Shutdown Sequence
     * 1. Set running_ = false (signals threads to exit)
     * 2. Write to event_fd (wakes listener from epoll_wait)
     * 3. Notify all workers (wakes them from condition_variable wait)
     * 4. Join listener thread (waits for it to finish)
     * 5. Join all worker threads (wait for them to finish)
     * 6. Close all client FDs
     * 7. Close epoll_fd and event_fd
     *
     * @note This method blocks until all threads have finished.
     *       In-flight requests will complete before workers exit.
     *
     * @warning Do not call stop() on a module that is not running.
     *          The behavior is undefined.
     */
    void stop() {
        running_ = false;

        // Wake up listener from epoll_wait
        uint64_t value = 1;
        write(event_fd_, &value, sizeof(value));

        // Wake up all workers from condition_variable wait
        queue_cv_.notify_all();

        // Wait for listener thread to finish
        if (listener_thread_.joinable()) {
            listener_thread_.join();
        }

        // Wait for all worker threads to finish
        for (auto& t : worker_threads_) {
            if (t.joinable()) {
                t.join();
            }
        }

        // Close all client FDs
        for (int fd : client_fds_) {
            close(fd);
        }

        // Close epoll and event FDs
        close(epoll_fd_);
        close(event_fd_);

        std::cout << "[GenAI] Module stopped\n";
    }

    /**
     * @brief Get the current size of the request queue
     *
     * Returns the number of requests currently waiting in the queue
     * to be processed by worker threads.
     *
     * @return Current queue size (number of pending requests)
     *
     * @par Thread Safety
     * This method is thread-safe and can be called concurrently
     * by multiple threads.
     *
     * @par Use Cases
     * - Monitoring: Track queue depth to detect backpressure
     * - Metrics: Collect statistics on request load
     * - Debugging: Verify queue is draining properly
     *
     * @note The queue size is momentary and may change immediately
     *       after this method returns.
     */
    size_t get_queue_size() const {
        std::lock_guard<std::mutex> lock(queue_mutex_);
        return request_queue_.size();
    }

private:
    /**
     * @brief Listener thread main loop
     *
     * Runs in a dedicated thread and monitors all registered client file
     * descriptors using epoll. When a client sends a request, this method
     * reads it, validates it, and enqueues it for worker processing.
     *
     * @par Event Loop
     * 1. Wait on epoll for events (timeout: 100ms)
     * 2. For each ready FD:
     *    - If event_fd: check for shutdown signal
     *    - If client FD: read request and enqueue
     * 3. If client disconnects: remove from epoll and close FD
     * 4. Loop until running_ is false
     *
     * @par Request Reading Flow
     * 1. Read RequestHeader (fixed size)
     * 2. Validate read succeeded (n > 0)
     * 3. Read input data (variable size based on header.input_size)
     * 4. Create Request structure
     * 5. Push to request_queue_
     * 6. Notify one worker via condition variable
     *
     * @param worker_id Unused parameter (for future per-worker stats)
     *
     * @note Runs with 100ms timeout on epoll_wait to periodically check
     *       the running_ flag for shutdown.
     */
    void listener_loop() {
        const int MAX_EVENTS = 64;  // Max events to process per epoll_wait
        struct epoll_event events[MAX_EVENTS];

        std::cout << "[GenAI] Listener thread started\n";

        while (running_) {
            // Wait for events on monitored FDs
            // Timeout of 100ms allows periodic check of running_ flag
            int nfds = epoll_wait(epoll_fd_, events, MAX_EVENTS, 100);

            if (nfds < 0 && errno != EINTR) {
                perror("epoll_wait");
                break;
            }

            // Process each ready FD
            for (int i = 0; i < nfds; i++) {
                // Check if this is the shutdown eventfd
                if (events[i].data.fd == event_fd_) {
                    // Shutdown signal - will exit loop when running_ is false
                    continue;
                }

                int client_fd = events[i].data.fd;

                // Read request header
                RequestHeader header;
                ssize_t n = read(client_fd, &header, sizeof(header));

                if (n <= 0) {
                    // Connection closed or error
                    std::cout << "[GenAI] Client fd " << client_fd << " disconnected\n";
                    epoll_ctl(epoll_fd_, EPOLL_CTL_DEL, client_fd, nullptr);
                    close(client_fd);
                    std::lock_guard<std::mutex> lock(clients_mutex_);
                    client_fds_.erase(client_fd);
                    continue;
                }

                // Read input data (may require multiple reads for large data)
                std::string input(header.input_size, '\0');
                size_t total_read = 0;
                while (total_read < header.input_size) {
                    ssize_t r = read(client_fd, &input[total_read], header.input_size - total_read);
                    if (r <= 0) break;
                    total_read += r;
                }

                // Create request and enqueue for processing
                Request req;
                req.client_fd = client_fd;
                req.request_id = header.request_id;
                req.operation = header.operation;
                req.input = std::move(input);

                {
                    // Critical section: modify request_queue_
                    std::lock_guard<std::mutex> lock(queue_mutex_);
                    request_queue_.push(std::move(req));
                }

                // Notify one worker thread that work is available
                queue_cv_.notify_one();

                std::cout << "[GenAI] Enqueued request " << header.request_id
                         << " from fd " << client_fd
                         << " (queue size: " << get_queue_size() << ")\n" << std::flush;
            }
        }

        std::cout << "[GenAI] Listener thread stopped\n";
    }

    /**
     * @brief Worker thread main loop
     *
     * Runs in each worker thread. Waits for requests to appear in the
     * queue, processes them (potentially blocking), and sends responses
     * back to clients.
     *
     * @par Worker Loop
     * 1. Wait on condition variable for queue to have work
     * 2. When notified, check if running_ is still true
     * 3. Pop request from queue (critical section)
     * 4. Process request (may block on I/O to external services)
     * 5. Send response back to client
     * 6. Loop back to step 1
     *
     * @par Request Processing
     * Currently simulates processing with a random sleep (100-500ms).
     * In production, this would call actual LLM APIs (OpenAI, local models, etc.).
     *
     * @par Response Sending
     * - Writes ResponseHeader first (fixed size)
     * - Writes output data second (variable size)
     * - Client reads both to get complete response
     *
     * @param worker_id Identifier for this worker (0 to num_workers_-1)
     *                  Used for logging and potentially for per-worker stats.
     *
     * @note Workers exit when running_ is set to false and queue is empty.
     *       In-flight requests will complete before workers exit.
     */
    void worker_loop(int worker_id) {
        std::cout << "[GenAI] Worker " << worker_id << " started\n";

        while (running_) {
            Request req;

            // Wait for work to appear in queue
            {
                std::unique_lock<std::mutex> lock(queue_mutex_);
                queue_cv_.wait(lock, [this] {
                    return !running_ || !request_queue_.empty();
                });

                if (!running_) break;

                if (request_queue_.empty()) continue;

                // Get request from front of queue
                req = std::move(request_queue_.front());
                request_queue_.pop();
            }

            // Simulate processing time (random sleep between 100-500ms)
            // In production, this would be actual LLM API calls
            unsigned int seed = req.request_id;
            int sleep_ms = 100 + (rand_r(&seed) % 400);  // 100-500ms

            std::cout << "[GenAI] Worker " << worker_id
                     << " processing request " << req.request_id
                     << " (sleep " << sleep_ms << "ms)\n";

            std::this_thread::sleep_for(std::chrono::milliseconds(sleep_ms));

            // Prepare response
            std::string output = "Processed: " + req.input;

            ResponseHeader resp;
            resp.request_id = req.request_id;
            resp.status_code = 0;
            resp.output_size = output.size();
            resp.processing_time_ms = sleep_ms;

            // Send response back to client
            write(req.client_fd, &resp, sizeof(resp));
            write(req.client_fd, output.data(), output.size());

            std::cout << "[GenAI] Worker " << worker_id
                     << " completed request " << req.request_id << "\n";
        }

        std::cout << "[GenAI] Worker " << worker_id << " stopped\n";
    }

    // ========================================================================
    // Member Variables
    // ========================================================================

    int num_workers_;             ///< Number of worker threads in the pool
    std::atomic<bool> running_;   ///< Flag indicating if module is running

    int epoll_fd_;                ///< epoll instance file descriptor
    int event_fd_;                ///< eventfd for shutdown notification

    std::thread listener_thread_;         ///< Thread that monitors client FDs
    std::vector<std::thread> worker_threads_;  ///< Thread pool for request processing

    std::queue<Request> request_queue_;   ///< FIFO queue of pending requests
    mutable std::mutex queue_mutex_;       ///< Protects request_queue_
    std::condition_variable queue_cv_;    ///< Notifies workers when queue has work

    std::unordered_set<int> client_fds_;  ///< Set of registered client FDs
    mutable std::mutex clients_mutex_;     ///< Protects client_fds_
};

// ============================================================================
// Client
// ============================================================================

/**
 * @class Client
 * @brief Simulates a ProxySQL thread (MySQL/PgSQL) making GenAI requests
 *
 * This class demonstrates how a client thread would interact with the
 * GenAI module in a real ProxySQL deployment. It creates a socketpair
 * connection, sends requests asynchronously, and polls for responses.
 *
 * @par Communication Pattern
 *
 * 1. Create socketpair() (2 FDs: read_fd and genai_fd)
 * 2. Pass genai_fd to GenAI module via register_client()
 * 3. Keep read_fd for monitoring responses
 * 4. Send requests via genai_fd (non-blocking)
 * 5. Poll read_fd for responses
 * 6. Process responses as they arrive
 *
 * @par Asynchronous Operation
 *
 * The key design principle is that the client never blocks waiting for
 * GenAI responses. Instead:
 * - Send requests and continue working
 * - Poll for responses periodically
 * - Handle responses when they arrive
 *
 * This allows the client thread to handle many concurrent requests
 * and continue serving other clients while GenAI processes.
 *
 * @par Usage Example
 * @code{.cpp}
 * // Create client
 * Client client("MySQL-Thread-1", 10);
 *
 * // Connect to GenAI module
 * client.connect_to_genai(genai_module);
 *
 * // Run (send requests, wait for responses)
 * client.run();
 *
 * // Clean up
 * client.close();
 * @endcode
 */
class Client {
public:

    /**
     * @brief Construct a Client with specified name and request count
     *
     * Creates a client instance that will send a specified number of
     * requests to the GenAI module.
     *
     * @param name Human-readable name for this client (e.g., "MySQL-Thread-1")
     * @param num_requests Number of requests to send (default: 5)
     *
     * @note The name is used for logging/debugging to identify which
     *       client is sending/receiving requests.
     */
    Client(const std::string& name, int num_requests = 5)
        : name_(name), num_requests_(num_requests), next_id_(1) {}

    /**
     * @brief Connect to the GenAI module
     *
     * Creates a socketpair and registers one end with the GenAI module.
     * The client keeps one end for reading responses.
     *
     * @param genai Reference to the GenAI module to connect to
     *
     * @pre GenAI module has been started (genai.start() was called)
     *
     * @post Socketpair is created
     * @post One end is registered with GenAI module
     * @post Other end is kept for reading responses
     * @post Both FDs are set to non-blocking mode
     *
     * @par Socketpair Creation
     * - socketpair() creates 2 connected Unix domain sockets
     * - fds[0] (read_fd_): Client reads responses from this
     * - fds[1] (genai_fd_): Passed to GenAI, GenAI writes responses to this
     * - Data written to one end can be read from the other end
     *
     * @par Connection Flow
     * 1. Create socketpair(AF_UNIX, SOCK_STREAM, 0, fds)
     * 2. Store fds[0] as read_fd_ (client reads responses here)
     * 3. Store fds[1] as genai_fd_ (pass to GenAI module)
     * 4. Call genai.register_client(genai_fd_) to register with module
     * 5. Set read_fd_ to non-blocking for polling
     *
     * @note This method is typically called once per client at initialization.
     */
    void connect_to_genai(GenAIModule& genai) {
        // Create socketpair for bidirectional communication
        int fds[2];
        if (socketpair(AF_UNIX, SOCK_STREAM, 0, fds) < 0) {
            perror("socketpair");
            exit(1);
        }

        read_fd_ = fds[0];   // Client reads responses from this
        genai_fd_ = fds[1];  // GenAI writes responses to this

        // Register write end with GenAI module
        genai.register_client(genai_fd_);

        // Set read end to non-blocking for async polling
        int flags = fcntl(read_fd_, F_GETFL, 0);
        fcntl(read_fd_, F_SETFL, flags | O_NONBLOCK);

        std::cout << "[" << name_ << "] Connected to GenAI (read_fd=" << read_fd_ << ")\n";
    }

    /**
     * @brief Send all requests and wait for all responses
     *
     * This method:
     * 1. Sends all requests immediately (non-blocking)
     * 2. Polls for responses periodically
     * 3. Processes responses as they arrive
     * 4. Returns when all responses have been received
     *
     * @pre connect_to_genai() has been called
     *
     * @post All requests have been sent
     * @post All responses have been received
     * @post pending_requests_ is empty
     *
     * @par Sending Phase
     * - Loop num_requests_ times
     * - Each iteration calls send_request()
     * - Requests are sent immediately, no waiting
     *
     * @par Receiving Phase
     * - Loop until completed_ == num_requests_
     * - Call process_responses() to check for new responses
     * - Sleep 50ms between checks (non-blocking poll)
     * - Process each response as it arrives
     *
     * @note In production, the 50ms sleep would be replaced by adding
     *       read_fd_ to the thread's epoll set along with other FDs.
     */
    void run() {
        // Send all requests immediately (non-blocking)
        for (int i = 0; i < num_requests_; i++) {
            send_request(i);
        }

        // Wait for all responses (simulate async handling)
        std::cout << "[" << name_ << "] Waiting for " << num_requests_ << " responses...\n";

        while (completed_ < num_requests_) {
            process_responses();
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
        }

        std::cout << "[" << name_ << "] All requests completed!\n";
    }

    /**
     * @brief Close the connection to GenAI module
     *
     * Closes both ends of the socketpair.
     *
     * @post read_fd_ is closed
     * @post genai_fd_ is closed
     * @post FDs are set to -1 (closed state)
     *
     * @note This should be called after run() completes.
     */
    void close() {
        if (read_fd_ >= 0) ::close(read_fd_);
        if (genai_fd_ >= 0) ::close(genai_fd_);
    }

private:

    /**
     * @brief Send a single request to the GenAI module
     *
     * Creates and sends a request with a generated input string.
     *
     * @param index Index of this request (used to generate input text)
     *
     * @par Request Creation
     * - Generate unique request_id (incrementing counter)
     * - Create input string: "name_ input #index"
     * - Fill in RequestHeader
     * - Write header to genai_fd_
     * - Write input data to genai_fd_
     * - Store request in pending_requests_ with timestamp
     *
     * @par Non-Blocking Operation
     * The write operations may block briefly, but in practice:
     * - Socket buffer is typically large enough
     * - If buffer is full, EAGAIN is returned (not handled in this demo)
     * - Client would need to retry in production
     *
     * @note This method increments next_id_ to ensure unique request IDs.
     */
    void send_request(int index) {
        // Create input string for this request
        std::string input = name_ + " input #" + std::to_string(index);
        uint64_t request_id = next_id_++;

        // Fill request header
        RequestHeader req;
        req.request_id = request_id;
        req.operation = OP_EMBEDDING;
        req.input_size = input.size();
        req.flags = 0;

        // Send request header
        write(genai_fd_, &req, sizeof(req));

        // Send input data
        write(genai_fd_, input.data(), input.size());

        // Track this request with timestamp for measuring round-trip time
        pending_requests_[request_id] = std::chrono::steady_clock::now();

        std::cout << "[" << name_ << "] Sent request " << request_id
                 << " (" << input << ")\n";
    }

    /**
     * @brief Check for and process any available responses
     *
     * Attempts to read a response from read_fd_. If a complete response
     * is available, processes it and updates tracking.
     *
     * @par Response Reading Flow
     * 1. Try to read ResponseHeader (non-blocking)
     * 2. If no data available (n <= 0), return immediately
     * 3. Read output data (may require multiple reads)
     * 4. Look up pending request by request_id
     * 5. Calculate round-trip time
     * 6. Log response details
     * 7. Remove from pending_requests_
     * 8. Increment completed_ counter
     *
     * @note This method is called periodically from run() to poll for
     *       responses. In production, read_fd_ would be in an epoll set.
     */
    void process_responses() {
        ResponseHeader resp;
        ssize_t n = read(read_fd_, &resp, sizeof(resp));

        if (n <= 0) {
            return;  // No data available yet
        }

        // Read output data
        std::string output(resp.output_size, '\0');
        size_t total_read = 0;
        while (total_read < resp.output_size) {
            ssize_t r = read(read_fd_, &output[total_read], resp.output_size - total_read);
            if (r <= 0) break;
            total_read += r;
        }

        // Find and process the matching pending request
        auto it = pending_requests_.find(resp.request_id);
        if (it != pending_requests_.end()) {
            auto start_time = it->second;
            auto end_time = std::chrono::steady_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(
                end_time - start_time).count();

            std::cout << "[" << name_ << "] Received response for request " << resp.request_id
                     << " (took " << duration << "ms, processed in "
                     << resp.processing_time_ms << "ms): " << output << "\n";

            pending_requests_.erase(it);
            completed_++;
        }
    }

    // ========================================================================
    // Member Variables
    // ========================================================================

    std::string name_;                       ///< Human-readable client identifier
    int num_requests_;                       ///< Total number of requests to send
    uint64_t next_id_;                       ///< Counter for generating unique request IDs
    int completed_ = 0;                      ///< Number of requests completed

    int read_fd_ = -1;                       ///< FD for reading responses from GenAI
    int genai_fd_ = -1;                      ///< FD for writing requests to GenAI

    /// Map of pending requests: request_id -> timestamp when sent
    std::unordered_map<uint64_t, std::chrono::steady_clock::time_point> pending_requests_;
};

// ============================================================================
// Main - Demonstration Entry Point
// ============================================================================

/**
 * @brief Main entry point for the GenAI module demonstration
 *
 * Creates a GenAI module with 4 workers and spawns 3 client threads
 * (simulating 2 MySQL threads and 1 PgSQL thread) that each send 3
 * concurrent requests.
 *
 * @par Execution Flow
 * 1. Create GenAI module with 4 workers
 * 2. Start the module (spawns listener and worker threads)
 * 3. Wait 100ms for module to initialize
 * 4. Create and start 3 client threads:
 *    - MySQL-Thread-1: 3 requests
 *    - MySQL-Thread-2: 3 requests
 *    - PgSQL-Thread-1: 3 requests
 * 5. Wait for all clients to complete
 * 6. Stop the GenAI module
 *
 * @par Expected Output
 * The program will output:
 * - Thread start/stop messages
 * - Client connection messages
 * - Request send/receive messages
 * - Timing information (round-trip time, processing time)
 * - Completion messages
 *
 * @return 0 on success, non-zero on error
 *
 * @note All clients are started with 50ms delays to demonstrate
 *       interleaved execution.
 */
int main() {
    std::cout << "=== GenAI Module Demonstration ===\n\n";

    // Create and start GenAI module with 4 worker threads
    GenAIModule genai(4);
    genai.start();

    std::this_thread::sleep_for(std::chrono::milliseconds(100));

    // Create multiple clients
    std::cout << "\n=== Creating Clients ===\n";

    std::vector<std::thread> client_threads;

    // Client 1: MySQL Thread simulation
    client_threads.emplace_back([&genai]() {
        Client client("MySQL-Thread-1", 3);
        client.connect_to_genai(genai);
        client.run();
        client.close();
    });

    std::this_thread::sleep_for(std::chrono::milliseconds(50));

    // Client 2: MySQL Thread simulation
    client_threads.emplace_back([&genai]() {
        Client client("MySQL-Thread-2", 3);
        client.connect_to_genai(genai);
        client.run();
        client.close();
    });

    std::this_thread::sleep_for(std::chrono::milliseconds(50));

    // Client 3: PgSQL Thread simulation
    client_threads.emplace_back([&genai]() {
        Client client("PgSQL-Thread-1", 3);
        client.connect_to_genai(genai);
        client.run();
        client.close();
    });

    // Wait for all clients to complete
    for (auto& t : client_threads) {
        if (t.joinable()) {
            t.join();
        }
    }

    std::cout << "\n=== All Clients Completed ===\n";

    // Stop GenAI module
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    genai.stop();

    std::cout << "\n=== Demonstration Complete ===\n";

    return 0;
}
