#include "molview/include/rdkit_bridge.hpp"

#include <GraphMol/FileParsers/MolSupplier.h>
#include <GraphMol/GraphMol.h>
#include <GraphMol/MolDraw2D/MolDraw2DCairo.h>
#include <GraphMol/SmilesParse/SmilesParse.h>

#include <algorithm>
#include <cstdint>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

namespace {

rust::Vec<std::uint8_t> render_molecules(
    const std::vector<std::unique_ptr<RDKit::ROMol>>& molecules
) {
    if (molecules.empty()) {
        throw std::invalid_argument("no valid molecules to render");
    }

    constexpr int panel_size = 300;
    const int columns = static_cast<int>(
        std::min<std::size_t>(molecules.size(), 3)
    );
    const int rows = static_cast<int>(
        (molecules.size() + columns - 1) / columns
    );

    std::vector<RDKit::ROMol*> molecule_views;
    molecule_views.reserve(molecules.size());
    for (const auto& molecule : molecules) {
        molecule_views.push_back(molecule.get());
    }

    RDKit::MolDraw2DCairo drawer{
        columns * panel_size,
        rows * panel_size,
        panel_size,
        panel_size,
    };
    drawer.drawMolecules(molecule_views);
    drawer.finishDrawing();

    const std::string png = drawer.getDrawingText();
    rust::Vec<std::uint8_t> result;
    result.reserve(png.size());
    for (const unsigned char byte : png) {
        result.push_back(byte);
    }
    return result;
}

}  // namespace

bool can_parse_smiles(rust::Str smiles) {
    const std::string input{
        smiles.data(),
        smiles.size(),
    };

    std::unique_ptr<RDKit::ROMol> molecule{
        RDKit::SmilesToMol(input),
    };

    return molecule != nullptr;
}

std::uint64_t count_sdf_molecules(
    rust::Slice<const std::uint8_t> sdf
) {
    if (sdf.empty()) {
        return 0;
    }

    const std::string contents{
        reinterpret_cast<const char*>(sdf.data()),
        sdf.size(),
    };

    std::istringstream input{
        contents,
        std::ios::in | std::ios::binary,
    };

    // Arguments:
    //   stream
    //   takeOwnership = false: `input` is stack-owned here
    //   sanitize = true
    //   removeHs = false
    //   strictParsing = true
    RDKit::ForwardSDMolSupplier supplier{
        &input,
        false,
        true,
        false,
        true,
    };

    std::uint64_t count = 0;

    while (!supplier.atEnd()) {
        std::unique_ptr<RDKit::ROMol> molecule{
            supplier.next(),
        };

        // RDKit may return nullptr for an invalid record while allowing
        // iteration to continue to later records.
        if (molecule) {
            ++count;
        }
    }

    return count;
}

rust::Vec<std::uint8_t> render_smiles_png(rust::Str smiles) {
    const std::string input{
        smiles.data(),
        smiles.size(),
    };

    std::unique_ptr<RDKit::ROMol> molecule{
        RDKit::SmilesToMol(input),
    };
    if (!molecule) {
        throw std::invalid_argument("invalid SMILES");
    }

    std::vector<std::unique_ptr<RDKit::ROMol>> molecules;
    molecules.push_back(std::move(molecule));
    return render_molecules(molecules);
}

rust::Vec<std::uint8_t> render_sdf_png(
    rust::Slice<const std::uint8_t> sdf
) {
    if (sdf.empty()) {
        throw std::invalid_argument("no valid molecules to render");
    }

    const std::string contents{
        reinterpret_cast<const char*>(sdf.data()),
        sdf.size(),
    };
    std::istringstream input{
        contents,
        std::ios::in | std::ios::binary,
    };
    RDKit::ForwardSDMolSupplier supplier{
        &input,
        false,
        true,
        false,
        true,
    };

    std::vector<std::unique_ptr<RDKit::ROMol>> molecules;
    while (!supplier.atEnd()) {
        std::unique_ptr<RDKit::ROMol> molecule{
            supplier.next(),
        };
        if (molecule) {
            molecules.push_back(std::move(molecule));
        }
    }

    return render_molecules(molecules);
}