#include "fastnet_interface.h"

#include <rtapi_ctype.h>
#include <rtapi.h>
#include <rtapi_app.h>
#include <hal.h>

#include <exception>
#include <string>

#include "metadata.pb.h"

MODULE_AUTHOR("zyp");
MODULE_DESCRIPTION("fastnet driver");
MODULE_LICENSE("BSD");

int comp_id;

class HalSharedMem {};

void* operator new(size_t size, const HalSharedMem) {
    return hal_malloc(size);
}

constexpr HalSharedMem hal_shared_mem;

class HalComponent {
    private:
        int id;
    
    public:
        HalComponent() {
            id = hal_init("fastnet");
        }

        ~HalComponent() {
            hal_exit(id);
        }

        int operator*() {
            return id;
        }

        void ready() {
            hal_ready(id);
        }
};

constexpr auto hal_type(hal_bit_t) { return HAL_BIT; }
constexpr auto hal_type(hal_float_t) { return HAL_FLOAT; }
constexpr auto hal_type(hal_s32_t) { return HAL_S32; }
constexpr auto hal_type(hal_u32_t) { return HAL_U32; }

template <typename T, hal_pin_dir_t Direction>
class HalPin {
    private:
        T* ptr;
    
    public:
        HalPin(HalComponent& comp, const std::string name) {
            hal_pin_new(name.c_str(), hal_type(T()), Direction, const_cast<void**>(reinterpret_cast<volatile void**>(&ptr)), *comp);
        }

        T& operator*() {
            return *ptr;
        }
};

using HalPinBitIn = HalPin<hal_bit_t, HAL_IN>;
using HalPinBitOut = HalPin<hal_bit_t, HAL_OUT>;
using HalPinU32In = HalPin<hal_u32_t, HAL_IN>;
using HalPinU32Out = HalPin<hal_u32_t, HAL_OUT>;
using HalPinS32In = HalPin<hal_s32_t, HAL_IN>;
using HalPinS32Out = HalPin<hal_s32_t, HAL_OUT>;

class FastnetSignal {
    public:
        virtual void update(std::span<uint8_t> buf) = 0;
};

class FastnetOutputBit : public FastnetSignal{
    private:
        HalPinBitIn* hal_pin;
        std::size_t offset;
    public:
        FastnetOutputBit(HalComponent& comp, const std::string& name, std::size_t offset) : offset(offset) {
            hal_pin = new (hal_shared_mem) HalPinBitIn(comp, name);
        }

        void update(std::span<uint8_t> buf) override {
            if (buf.size() <= offset / 8) {
                return;
            }

            if (**hal_pin) {
                buf[offset / 8] |= 1 << (offset % 8);
            }
        }
};

class FastnetOutputU32 : public FastnetSignal{
    private:
        HalPinU32In* hal_pin;
        std::size_t offset;
    public:
        FastnetOutputU32(HalComponent& comp, const std::string& name, std::size_t offset) : offset(offset) {
            hal_pin = new (hal_shared_mem) HalPinU32In(comp, name);
        }

        void update(std::span<uint8_t> buf) override {
            if (buf.size() <= offset / 8 + 3) {
                return;
            }

            uint32_t value = **hal_pin;

            buf[offset / 8 + 0] = value >> 0;
            buf[offset / 8 + 1] = value >> 8;
            buf[offset / 8 + 2] = value >> 16;
            buf[offset / 8 + 3] = value >> 24;
        }
};

class FastnetOutputS32 : public FastnetSignal{
    private:
        HalPinS32In* hal_pin;
        std::size_t offset;
    public:
        FastnetOutputS32(HalComponent& comp, const std::string& name, std::size_t offset) : offset(offset) {
            hal_pin = new (hal_shared_mem) HalPinS32In(comp, name);
        }

        void update(std::span<uint8_t> buf) override {
            if (buf.size() <= offset / 8 + 3) {
                return;
            }

            uint32_t value = **hal_pin;

            buf[offset / 8 + 0] = value >> 0;
            buf[offset / 8 + 1] = value >> 8;
            buf[offset / 8 + 2] = value >> 16;
            buf[offset / 8 + 3] = value >> 24;
        }
};

class FastnetInputBit : public FastnetSignal {
    private:
        HalPinBitOut* hal_pin;
        std::size_t offset;
    public:
        FastnetInputBit(HalComponent& comp, const std::string& name, std::size_t offset) : offset(offset) {
            hal_pin = new (hal_shared_mem) HalPinBitOut(comp, name);
        }

        void update(std::span<uint8_t> buf) override {
            if (buf.size() <= offset / 8) {
                return;
            }

            **hal_pin = buf[offset / 8] & 1 << (offset % 8) ? 1 : 0;
        }
};

class FastnetInputU32 : public FastnetSignal{
    private:
        HalPinU32Out* hal_pin;
        std::size_t offset;
    public:
        FastnetInputU32(HalComponent& comp, const std::string& name, std::size_t offset) : offset(offset) {
            hal_pin = new (hal_shared_mem) HalPinU32Out(comp, name);
        }

        void update(std::span<uint8_t> buf) override {
            if (buf.size() <= offset / 8 + 3) {
                return;
            }

            uint32_t value = 0;

            value |= buf[offset / 8 + 0] << 0;
            value |= buf[offset / 8 + 1] << 8;
            value |= buf[offset / 8 + 2] << 16;
            value |= buf[offset / 8 + 3] << 24;

            **hal_pin = value;
        }
};

class FastnetInputS32 : public FastnetSignal{
    private:
        HalPinS32Out* hal_pin;
        std::size_t offset;
    public:
        FastnetInputS32(HalComponent& comp, const std::string& name, std::size_t offset) : offset(offset) {
            hal_pin = new (hal_shared_mem) HalPinS32Out(comp, name);
        }

        void update(std::span<uint8_t> buf) override {
            if (buf.size() <= offset / 8 + 3) {
                return;
            }

            uint32_t value = 0;

            value |= buf[offset / 8 + 0] << 0;
            value |= buf[offset / 8 + 1] << 8;
            value |= buf[offset / 8 + 2] << 16;
            value |= buf[offset / 8 + 3] << 24;

            **hal_pin = value;
        }
};

struct FastnetDevice {
    HalComponent& comp;
    uint8_t addr;
    std::vector<FastnetSignal*> outputs;
    std::vector<FastnetSignal*> inputs;
    std::size_t data_size;
    std::size_t data_offset;

    void update_outputs(std::span<uint8_t> buf) {
        for (auto output : outputs) {
            output->update(buf);
        }
    }

    void update_inputs(std::span<uint8_t> buf) {
        for (auto input : inputs) {
            input->update(buf);
        }
    }

    void process_metadata(std::string prefix, const google::protobuf::RepeatedPtrField<Group::Item>& items) {
        for(auto& item : items) {
            auto name = prefix + "." + item.name();

            if (item.has_group()) {
                process_metadata(name, item.group().items());

            } else if(item.has_signal() and item.signal().direction() == OUTPUT) {
                switch(item.signal().type()) {
                    case BIT:
                        outputs.push_back(new FastnetOutputBit(comp, name, item.signal().offset()));
                        break;

                    case U32:
                        outputs.push_back(new FastnetOutputU32(comp, name, item.signal().offset()));
                        break;

                    case S32:
                        outputs.push_back(new FastnetOutputS32(comp, name, item.signal().offset()));
                        break;

                }
            } else if(item.has_signal() and item.signal().direction() == INPUT) {
                switch(item.signal().type()) {
                    case BIT:
                        inputs.push_back(new FastnetInputBit(comp, name, item.signal().offset()));
                        break;

                    case U32:
                        inputs.push_back(new FastnetInputU32(comp, name, item.signal().offset()));
                        break;

                    case S32:
                        inputs.push_back(new FastnetInputS32(comp, name, item.signal().offset()));
                        break;

                }
            }

            if(item.has_signal()) {
                data_size = std::max<std::size_t>(data_size, (item.signal().offset() + 32) / 32 * 4);
            }
        }
    }

};

class Module {
    private:
        HalComponent comp;
        FastnetInterface* interface;

        std::vector<FastnetDevice> devices;

        std::size_t data_packet_size;

    public:
        Module() {
            rtapi_set_msg_level(RTAPI_MSG_INFO);

            interface = new FastnetInterface("enp3s0");

            rtapi_print_msg(RTAPI_MSG_INFO, "Enumerating chain.\n");

            std::size_t count = interface->enumerate();

            rtapi_print_msg(RTAPI_MSG_INFO, "Found %d devices.\n", count);

            for (std::size_t addr = 0; addr < count; addr++) {
                auto& dev = devices.emplace_back(comp, addr);

                std::array<uint8_t, 1024> buf;
                auto res = interface->get_metadata(addr, buf);

                DeviceInfo device_info;
                device_info.ParseFromArray(res.data(), res.size());

                dev.process_metadata("fastnet." + std::to_string(addr), device_info.signals());
            }

            data_packet_size = 4;
            for (auto& dev : devices) {
                dev.data_offset = data_packet_size;
                data_packet_size += 4 + dev.data_size;
            }

            hal_export_funct("fastnet.update", [](void* arg, long period) {
                reinterpret_cast<Module*>(arg)->update(period);
            }, this, true, false, *comp);

            comp.ready();
        }

        void update(long period) {
            uint8_t rawbuf[data_packet_size] = {
                3, 0, 0, 0, // Command header
            };

            std::span<uint8_t> buf {rawbuf, data_packet_size};

            //std::array<uint8_t, 20> buf = {
            //    3, 0, 0, 0, // Command header
            //    0, 0, 4, 0, // Device 0 header
            //    0, 0, 0, 0, // Device 0 data
            //    1, 0, 4, 0, // Device 1 header
            //    0, 0, 0, 0, // Device 1 data
            //};

            for (auto& dev : devices) {
                //dev.update_outputs(std::span(buf).subspan(8 + 8 * dev.addr, 4));
                buf[dev.data_offset] = dev.addr;
                buf[dev.data_offset + 2] = dev.data_size;
                dev.update_outputs(buf.subspan(dev.data_offset + 4, dev.data_size));
            }

            interface->send_cmd(buf, buf);

            for (auto& dev : devices) {
                //dev.update_inputs(std::span(buf).subspan(8 + 8 * dev.addr, 4));
                dev.update_inputs(buf.subspan(dev.data_offset + 4, dev.data_size));
            }
        }
};

Module* m;

extern "C" int rtapi_app_main() {
    try {
        m = new Module;
    } catch(std::exception& e) {
        rtapi_print_msg(RTAPI_MSG_ERR, "fastnet ERROR: %s\n", e.what());
        return -EFAULT;
    }

    return 0;
}

extern "C" void rtapi_app_exit() {
	if(m) {
        delete m;
    }
}