#include "../../include/dv-sdk/module.hpp"

#include <boost/filesystem.hpp>
#include <opencv2/opencv.hpp>

#include <iomanip>
#include <sstream>
#include <string>

class ImageOutput : public dv::ModuleBase {
private:
	boost::filesystem::path path;
	std::string prefix;
	int32_t counter;
	int64_t generated;

public:
	static void initInputs(dv::InputDefinitionList &in) {
		in.addFrameInput("frames");
	}

	static const char *initDescription() {
		return ("Write out frames as PNG images.");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		config.add("directory", dv::ConfigOption::directoryOption("Directory to write images in."));
		config.add("imageFilePrefix", dv::ConfigOption::stringOption("Prefix to generate a file name.", "dv_frame"));
		config.add("imageFileCounter",
			dv::ConfigOption::intOption("Start number for file name suffix counter.", 1, 0, INT32_MAX));

		config.add("generatedImagesCount", dv::ConfigOption::statisticOption("Number of images written so far."));

		config.setPriorityOptions({"directory", "imageFilePrefix", "imageFileCounter"});
	}

	ImageOutput() {
		path = boost::filesystem::path(config.getString("directory"));

		if (path.empty()) {
			throw std::runtime_error("Output directory not specified.");
		}

		if (!boost::filesystem::exists(path) || !boost::filesystem::is_directory(path)) {
			throw std::runtime_error("Output directory doesn't exist or is not a directory.");
		}

		prefix  = config.getString("imageFilePrefix") + "_";
		counter = config.getInt("imageFileCounter");

		generated = 0;
	}

	void run() override {
		auto frame = inputs.getFrameInput("frames").frame();

		std::stringstream ss;
		ss << std::setw(10) << std::setfill('0') << counter;

		std::string fileName = prefix + ss.str() + ".png";

		boost::filesystem::path filePath = path / fileName;

		cv::imwrite(filePath.string(), *frame.getMatPointer());

		counter++;

		config.setLong("generatedImagesCount", ++generated);
	}

	~ImageOutput() override {
		config.setInt("imageFileCounter", counter);

		config.set<dv::CfgType::LONG>("generatedImagesCount", generated, true);

		log.info << generated << " frames saved in: " << path << std::endl;
	}
};

registerModuleClass(ImageOutput)
