/*
 * Copyright (C) 2018 Fabio D'Urso <durso@dmi.unict.it>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
*/
#include "CameraBufferPlugin.hh"

#include <gazebo/rendering/rendering.hh>

#include <stdio.h>
#include <pthread.h>
#include <unistd.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>

using namespace gazebo;

GZ_REGISTER_SENSOR_PLUGIN(GzUav::CameraBufferPlugin)

static void SendAll(int fd, const uint8_t *data, size_t len)
{
    while (len != 0)
    {
        int r = send(fd, data, len, MSG_NOSIGNAL);
        if (r < 0)
            return; // failure

        data += r;
        len -= r;
    }
}

GzUav::CameraBufferPlugin::CameraBufferPlugin()
: serv(-1)
{
    pthread_mutex_init(&this->lastDataMutex, nullptr);
}

GzUav::CameraBufferPlugin::~CameraBufferPlugin()
{
    // If the server thread was started, wait for it to stop before destroying
    // this object
    if (this->serv != -1)
    {
        pthread_cancel(this->serverThread);
        pthread_join(this->serverThread, nullptr);
    }

    pthread_mutex_destroy(&this->lastDataMutex);
}

void GzUav::CameraBufferPlugin::Load(sensors::SensorPtr _sensor, sdf::ElementPtr _sdf)
{
    // Call base class
    CameraPlugin::Load(_sensor, _sdf);

    // Extract the root model's name from the fully scoped name of the sensor
    // e.g. worldname::rootmodelname::submodelname::linkname::sensorname
    std::string scopedName = _sensor->ScopedName();
    int modelBegin = scopedName.find("::");
    int modelEnd = scopedName.find("::", modelBegin + 2);
    std::string modelName(&scopedName[modelBegin + 2], &scopedName[modelEnd]);

    // Start TCP server
    struct sockaddr_in addr;
    int opt_val = 1, port = atoi(getenv(("GZUAV_CAMBUFFER_PORT-" + modelName).c_str()));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = htonl(INADDR_ANY);

    this->serv = socket(AF_INET, SOCK_STREAM, 0);
    setsockopt(this->serv, SOL_SOCKET, SO_REUSEADDR, &opt_val, sizeof(int));
    if (bind(this->serv, (struct sockaddr*)&addr, sizeof(addr)) < 0)
        gzthrow("[CameraBufferPlugin] bind failed");

    listen(this->serv, 1);
    this->sock = -1;

    // Start server thread
    pthread_create(&this->serverThread, nullptr, ThreadFunc, (void*)this);
}

void GzUav::CameraBufferPlugin::OnNewFrame(const unsigned char *_image,
                unsigned int _width, unsigned int _height,
                unsigned int _depth, const std::string &_format)
{
    common::Time ts = this->camera->GetScene()->SimTime();
    size_t bytes = this->camera->ImageByteSize();

    pthread_mutex_lock(&this->lastDataMutex);
    this->lastDataTimestamp = ts;
    this->lastData.assign(_image, _image + bytes);
    pthread_mutex_unlock(&this->lastDataMutex);
}

void GzUav::CameraBufferPlugin::ServerThread()
{
    // Wait for the client
    while (this->sock < 0)
    {
        struct sockaddr_in addr;
        socklen_t len = sizeof(addr);
        this->sock = accept(this->serv, (struct sockaddr*)&addr, &len);
    }

    // Stop accepting connections
    close(this->serv);

    // Send image description
    char imageDescr[128 + this->format.length()];
    sprintf(imageDescr, "%d %d %d %s", this->width, this->height, this->depth,
            this->format.c_str());
    SendBlock(imageDescr, strlen(imageDescr));

    // Reply to image requests
    while (true)
    {
        int dummy;
        if (recv(this->sock, &dummy, 1, 0) != 1)
            return;

        pthread_setcancelstate(PTHREAD_CANCEL_DISABLE, &dummy);

        // This block cannot be canceled by pthread_cancel
        {
            // Take ownership of the current buffer
            std::vector<unsigned char> data;
            pthread_mutex_lock(&this->lastDataMutex);
            double ts = this->lastDataTimestamp.Double();
            data.swap(this->lastData);
            pthread_mutex_unlock(&this->lastDataMutex);

            if (data.size() != 0)
            {
                // Append timestamp to image data
                unsigned char *tsBytes = (unsigned char*)&ts;
                data.insert(data.end(), tsBytes, tsBytes + sizeof(ts));

                SendBlock(data.data(), data.size());
            }
            else
            {
                // No image is available, send empty response
                SendBlock(nullptr, 0);
            }
        }

        pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, &dummy);
    }
}

void *GzUav::CameraBufferPlugin::ThreadFunc(void *me)
{
    ((CameraBufferPlugin*)me)->ServerThread();
    return nullptr;
}

void GzUav::CameraBufferPlugin::SendBlock(const void *data, size_t len)
{
    // Enable TCP_CORK
    int opt_val = 1;
    setsockopt(this->sock, IPPROTO_TCP, TCP_CORK, &opt_val, sizeof(opt_val));

    // Send data length as a uint64_t followed by actual data
    uint64_t len64 = len;
    SendAll(this->sock, (const uint8_t*)&len64, sizeof(uint64_t));
    if (len > 0)
        SendAll(this->sock, (const uint8_t*)data, len);

    // Disable TCP_CORK (this has the effect of immediatly flushing data)
    opt_val = 0;
    setsockopt(this->sock, IPPROTO_TCP, TCP_CORK, &opt_val, sizeof(opt_val));
}
