#include "ccsds_apid_demux.h"
#include "bit_container.h"
#include "common/ccsds/ccsds_aos/demuxer.h"
#include "core/style.h"
#include "imgui/imgui.h"
#include "logger.h"
#include <fstream>
#include <map>
#include <string>

namespace satdump
{
    std::string CCSDSAPIDDemuxTool::getName() { return "CCSDS APID Demux"; }

    void CCSDSAPIDDemuxTool::renderMenu(std::shared_ptr<BitContainer> &container, bool is_busy)
    {
        if (is_busy)
            style::beginDisabled();

        if (ImGui::RadioButton("AOS", aos_mode == 1))
            aos_mode = 1;
        if (ImGui::RadioButton("TM", aos_mode == 0))
            aos_mode = 0;

        ImGui::InputInt("CADU Width (bytes)", &cadu_size_bytes);
        ImGui::InputInt("MPDU Data Size", &mpdu_data_size);
        ImGui::InputInt("MPDU Insert Zone Size", &insert_zone_size);
        ImGui::Checkbox("Split APIDs", &split_apid);

        if (ImGui::Button("Perform###2"))
            should_process = true;

        if (is_busy)
            style::endDisabled();
    }

    bool CCSDSAPIDDemuxTool::needToProcess() { return should_process; }

    void CCSDSAPIDDemuxTool::setProcessed() { should_process = false; }

    void CCSDSAPIDDemuxTool::process(std::shared_ptr<BitContainer> &container, float &process_progress)
    {
        logger->critical("Demux!!!!");

        uint8_t *ptr = container->get_ptr();
        size_t size = container->get_ptr_size();

        if (split_apid)
        {
            struct APIDOut
            {
                APIDOut(std::string f) : tmpfile(f), fileout(f, std::ios::binary) {}

                std::string tmpfile;
                std::ofstream fileout;
                std::vector<BitContainer::FrameDef> frms;
                size_t pos = 0;
            };
            std::map<int, std::shared_ptr<APIDOut>> apids_outs;

            ccsds::ccsds_aos::Demuxer demuxer_vcid_aos(mpdu_data_size, insert_zone_size, insert_zone_size);
            ccsds::ccsds_aos::Demuxer demuxer_vcid_tm(mpdu_data_size, insert_zone_size, insert_zone_size);

            for (int i = 0; i < size; i += cadu_size_bytes)
            {
                auto frm = aos_mode ? demuxer_vcid_aos.work(ptr + i) : demuxer_vcid_tm.work(ptr + i);

                for (auto &f : frm)
                {
                    if (apids_outs.count(f.header.apid) == 0)
                    {
                        char name[1000];
                        tmpnam(name);
                        std::string tmpfile = name;
                        apids_outs.emplace(f.header.apid, std::make_shared<APIDOut>(tmpfile));
                    }

                    auto &v = apids_outs[f.header.apid];

                    v->fileout.write((char *)f.header.raw, 6);
                    v->fileout.write((char *)f.payload.data(), f.payload.size());

                    v->frms.push_back({v->pos * 8, (f.payload.size() + 6) * 8});
                    v->pos += f.payload.size() + 6;
                }

                process_progress = double(i) / double(size);
            }

            for (auto &v : apids_outs)
            {
                v.second->fileout.close();

                if (v.second->frms.size())
                {
                    std::shared_ptr<satdump::BitContainer> newbitc = std::make_shared<satdump::BitContainer>("APID " + std::to_string(v.first), v.second->tmpfile, v.second->frms);
                    newbitc->d_bitperiod = 8192;
                    newbitc->init_display();
                    newbitc->d_is_temporary = true;

                    if (container->bitview != nullptr)
                        ((BitViewHandler *)container->bitview)->addSubHandler(std::make_shared<BitViewHandler>(newbitc));
                    else
                        logger->error("Can't add container!");
                }
            }
        }
        else
        {
            char name[1000];
            tmpnam(name);
            std::string tmpfile = name;
            std::ofstream fileout = std::ofstream(tmpfile, std::ios::binary);
            std::vector<BitContainer::FrameDef> frms;

            ccsds::ccsds_aos::Demuxer demuxer_vcid_aos(mpdu_data_size, insert_zone_size, insert_zone_size);
            ccsds::ccsds_aos::Demuxer demuxer_vcid_tm(mpdu_data_size, insert_zone_size, insert_zone_size);

            size_t pos = 0;
            for (int i = 0; i < size; i += cadu_size_bytes)
            {
                auto frm = aos_mode ? demuxer_vcid_aos.work(ptr + i) : demuxer_vcid_tm.work(ptr + i);

                for (auto &f : frm)
                {
                    fileout.write((char *)f.header.raw, 6);
                    fileout.write((char *)f.payload.data(), f.payload.size());

                    frms.push_back({pos * 8, (f.payload.size() + 6) * 8});
                    pos += f.payload.size() + 6;
                }

                process_progress = double(i) / double(size);
            }

            fileout.close();

            if (frms.size())
            {
                std::shared_ptr<satdump::BitContainer> newbitc = std::make_shared<satdump::BitContainer>("APIDs", tmpfile, frms);
                newbitc->d_bitperiod = 8192;
                newbitc->init_display();
                newbitc->d_is_temporary = true;

                if (container->bitview != nullptr)
                    ((BitViewHandler *)container->bitview)->addSubHandler(std::make_shared<BitViewHandler>(newbitc));
                else
                    logger->error("Can't add container!");
            }
        }
    }

} // namespace satdump