#include "udp_client.h"
#include <cstring>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include "esp_log.h"

static const char* TAG = "UDPClient";

UDPClient::UDPClient()
    : sock_fd_(-1)
    , remote_port_(0)
    , configured_(false)
    , initialized_(false) {
    memset(&remote_addr_, 0, sizeof(remote_addr_));
}

UDPClient::~UDPClient() {
    close();
}

esp_err_t UDPClient::init() {
    if (initialized_) {
        ESP_LOGW(TAG, "Already initialized");
        return ESP_OK;
    }

    sock_fd_ = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock_fd_ < 0) {
        ESP_LOGE(TAG, "Failed to create socket: errno %d", errno);
        return ESP_FAIL;
    }

    // Set socket to non-blocking mode
    esp_err_t err = set_nonblocking();
    if (err != ESP_OK) {
        ::close(sock_fd_);
        sock_fd_ = -1;
        return err;
    }

    initialized_ = true;
    ESP_LOGI(TAG, "UDP client initialized (fd=%d)", sock_fd_);
    return ESP_OK;
}

esp_err_t UDPClient::set_nonblocking() {
    int flags = fcntl(sock_fd_, F_GETFL, 0);
    if (flags < 0) {
        ESP_LOGE(TAG, "Failed to get socket flags: errno %d", errno);
        return ESP_FAIL;
    }

    if (fcntl(sock_fd_, F_SETFL, flags | O_NONBLOCK) < 0) {
        ESP_LOGE(TAG, "Failed to set non-blocking mode: errno %d", errno);
        return ESP_FAIL;
    }

    return ESP_OK;
}

esp_err_t UDPClient::configure(const std::string& ip, uint16_t port) {
    if (ip.empty() || port == 0) {
        ESP_LOGE(TAG, "Invalid IP or port");
        return ESP_ERR_INVALID_ARG;
    }

    struct in_addr addr;
    if (inet_pton(AF_INET, ip.c_str(), &addr) != 1) {
        ESP_LOGE(TAG, "Invalid IP address format: %s", ip.c_str());
        return ESP_ERR_INVALID_ARG;
    }

    remote_addr_.sin_family = AF_INET;
    remote_addr_.sin_port = htons(port);
    remote_addr_.sin_addr = addr;

    remote_ip_ = ip;
    remote_port_ = port;
    configured_ = true;

    ESP_LOGI(TAG, "Configured endpoint: %s:%u", ip.c_str(), port);
    return ESP_OK;
}

esp_err_t UDPClient::send_command(const std::string& command) {
    if (!initialized_) {
        ESP_LOGE(TAG, "Not initialized");
        return ESP_FAIL;
    }

    if (!configured_) {
        ESP_LOGE(TAG, "Endpoint not configured");
        return ESP_FAIL;
    }

    ssize_t sent = sendto(sock_fd_, command.c_str(), command.length(), 0,
                          (struct sockaddr*)&remote_addr_, sizeof(remote_addr_));

    if (sent < 0) {
        ESP_LOGE(TAG, "Failed to send command '%s': errno %d", command.c_str(), errno);
        return ESP_FAIL;
    }

    ESP_LOGD(TAG, "Sent command: %s (%d bytes)", command.c_str(), (int)sent);
    return ESP_OK;
}

esp_err_t UDPClient::receive_response(std::string& response, int timeout_ms) {
    if (!initialized_) {
        ESP_LOGE(TAG, "Not initialized");
        return ESP_FAIL;
    }

    // Setup select() for timeout
    fd_set read_fds;
    FD_ZERO(&read_fds);
    FD_SET(sock_fd_, &read_fds);

    struct timeval timeout;
    struct timeval* timeout_ptr = nullptr;

    if (timeout_ms >= 0) {
        timeout.tv_sec = timeout_ms / 1000;
        timeout.tv_usec = (timeout_ms % 1000) * 1000;
        timeout_ptr = &timeout;
    }

    int ret = select(sock_fd_ + 1, &read_fds, nullptr, nullptr, timeout_ptr);

    if (ret < 0) {
        ESP_LOGE(TAG, "select() failed: errno %d", errno);
        return ESP_FAIL;
    }

    if (ret == 0) {
        ESP_LOGD(TAG, "Receive timeout (%d ms)", timeout_ms);
        return ESP_ERR_TIMEOUT;
    }

    // Data is available
    char buffer[512];
    struct sockaddr_in from_addr;
    socklen_t from_len = sizeof(from_addr);

    ssize_t received = recvfrom(sock_fd_, buffer, sizeof(buffer) - 1, 0,
                                (struct sockaddr*)&from_addr, &from_len);

    if (received < 0) {
        ESP_LOGE(TAG, "recvfrom() failed: errno %d", errno);
        return ESP_FAIL;
    }

    buffer[received] = '\0';
    response = std::string(buffer, received);

    ESP_LOGD(TAG, "Received response: %s (%d bytes)", response.c_str(), (int)received);
    return ESP_OK;
}

void UDPClient::close() {
    if (sock_fd_ >= 0) {
        ::close(sock_fd_);
        ESP_LOGI(TAG, "Socket closed");
        sock_fd_ = -1;
    }

    initialized_ = false;
    configured_ = false;
    remote_ip_.clear();
    remote_port_ = 0;
}