#include "image.h"
#include "map.h"

#include "contrib_cmaps.h"
#include "turbo_colormap.hpp"
#include "vnum.h"

#include <iostream>

#include <kblib/algorithm.h>
#include <kblib/random.h>

#include "tclap/CmdLine.h"

using namespace std::literals;

template <typename T>
struct arg_pair {
 private:
	T a_{};
	std::optional<T> b_{};

 public:
	[[nodiscard]] auto a() const noexcept -> T { return a_; }
	[[nodiscard]] auto b() const noexcept -> T { return b_.value_or(a_); }

	friend auto operator>>(std::istream& in, arg_pair<T>& p) -> std::istream& {
		in >> p.a_;
		if (! in.fail()) {
			auto c = in.peek();
			if (c != kblib::eof<>) {
				if (",;:"sv.find(static_cast<char>(c)) != std::string::npos) {
					in.ignore();
				}
				in.clear();
				p.b_.emplace();
				in >> *p.b_;
			} else {
				p.b_.reset();
				in.clear();
			}
		}
		return in;
	}
};

template <typename NoiseEngine, typename OutputOperation,
          typename... NoiseParams, typename... OutputParams>
void operate(NoiseEngine noise, std::tuple<NoiseParams...> noise_params,
             OutputOperation out, std::tuple<OutputParams...> output_params) {
	std::apply(out, std::apply(noise, noise_params), output_params);
}

namespace TCLAP {
template <typename T>
struct ArgTraits<arg_pair<T>> {
	using ValueCategory = ValueLike;
};
} // namespace TCLAP

class SomeOf : public TCLAP::AnyOf {
	using AnyOf::AnyOf;
	auto validate() -> bool override {
		return std::none_of(begin(), end(), std::mem_fn(&TCLAP::Arg::isSet));
	}
	[[nodiscard]] auto isRequired() const -> bool override { return true; }
};

cmap_t read_cmap_from_file(std::string filename) {
	std::vector<sRGB_tuple> data;
	auto contents
	    = kblib::split_dsv(kblib::try_get_file_contents(filename), '\n');
	data.reserve(contents.size());
	for (const auto& triple : contents) {
		auto& d = data.emplace_back();
		auto e = std::sscanf(triple.c_str(), "%lf, %lf, %lf", &d.red(),
		                     &d.green(), &d.blue());
		if (e != 3) {
			throw std::invalid_argument("invalid triple");
		}
	}
	return cmap_t{1, {"in", sRGB_to_cmap(data)}};
}

void write_cmap(const cmap_t& cmap) {
	for (const auto i : kblib::range(cmap.size())) {
		const auto& [name, pal] = cmap[i];
		std::cout << ":" << name << ':' << pal.size() << '\n';
		for (const auto j : kblib::range(pal.size())) {
			std::cout << color_to_hex(pal[j]) << (j % 8 == 7 ? '\n' : ' ');
		}
		std::cout << '\n';
	}
}

auto draw_cmaps(std::string filename, std::string format,
                std::string_view cmapfilename) {
	const auto width = *kblib::max_element(
	    kblib::transformer(
	        begin(current_cmap),
	        +[](const cmap_t::value_type& pal) { return pal.second.size(); }),
	    end(current_cmap), std::less<>{});
	pnm::image<pnm::rgb_pixel> img(width, current_cmap.size());
	for (const auto cmap_n : kblib::range(current_cmap.size())) {
		const auto& ccmap = current_cmap[cmap_n].second;
		for (const auto& color_n : kblib::range(width)) {
			auto v = color_n * ccmap.size() / width;
			img[cmap_n][color_n]
			    = {ccmap[v].red(), ccmap[v].green(), ccmap[v].blue()};
		}
	}
	std::string comment
	    = kblib::concat("Color palette generated by MapGen v", vnum);
	if (not cmapfilename.empty()) {
		comment += kblib::concat(" from ", cmapfilename);
	}
	comment += '.';
	auto ec = pnm::finalize_to_png(img, filename, format, comment);
	if (ec) {
		std::cerr << "Error writing file.\n";
		return ec.value();
	}
	return 0;
}

auto main(int argc, char** argv) -> int {
	try {
		TCLAP::CmdLine cmd("", ' ', vnum);

		TCLAP::ValueArg<std::string> outfile(
		    "o", "output-file", "The file to write to (- for stdout)", false, "",
		    "FILENAME");
		TCLAP::ValueArg<std::string> format(
		    "f", "format",
		    "The file format to save to (automatically determined from extension "
		    "if not specified, falls back to PNG), applies to -o and -s",
		    false, "", "EXT", cmd);

		TCLAP::ValueArg<arg_pair<float>> corner(
		    "c", "corner",
		    "The coordinates of the top-left coordinate of the area (if only one "
		    "value is provided, it is used for both coordinates)",
		    false, {}, "LEFT[,UPPER]", cmd);

		TCLAP::ValueArg<float> width("w", "width", "The width of the output area",
		                             false, 0, "WIDTH", cmd);

		TCLAP::ValueArg<std::uint16_t> res(
		    "r", "res", "The resolution of the image", false, 256, "X", cmd);

		TCLAP::ValueArg<std::string> cmap("m", "cmap", "The colormap ID to use",
		                                  false, "1", "CMAP", cmd);

		TCLAP::ValueArg<std::string> cmapfile("M", "cmapfile",
		                                      "The cmap file to load cmaps from",
		                                      false, "", "FILE.pc", cmd);
		TCLAP::ValueArg<std::string> paramsfile(
		    "p", "params-file", "the file to read generation parameters from",
		    false, "", "FILE.dat", cmd);
		TCLAP::ValueArg<std::string> recreate(
		         "R", "recreate", "recreate a rendered image from its comment",
		         false, "", "IMAGE_FILE", cmd);

		TCLAP::MultiSwitchArg list_cmap(
		    "C", "list-cmap",
		    "List the names and lengths of all available cmaps. If given twice, "
		    "print their hex values too.");
		TCLAP::MultiSwitchArg list_params(
		    "P", "dump-params",
		    "print all generation parameters to stdout. If given twice, print a "
		    "description of each value too.");
		TCLAP::ValueArg<std::string> save_cmap("s", "save-cmaps",
		                                       "Save all cmaps to image", false,
		                                       "cmaps.png", "IMAGE_FILE");
		TCLAP::ValueArg<std::string> convert_srgb(
		    "", "convert", "A CSV file of sRGB tuples, one per line", false, "",
		    "srgb_file");

		SomeOf output(cmd);
		output.add(outfile);
		output.add(list_cmap);
		output.add(list_params);
		output.add(save_cmap);
		output.add(convert_srgb);

		cmd.parse(argc, argv);

#if 0
		auto rng = kblib::seeded<kblib::best_lcgs::lcg32>();
		auto dist = std::uniform_int_distribution<std::uint8_t>();
		pnm::image<pnm::rgb_pixel> img(res.getValue().a(), res.getValue().b());

		for (auto x : kblib::range(img.width())) {
			for (auto y : kblib::range(img.height())) {
				img[x][y] = {dist(rng), dist(rng), dist(rng)};
			}
		}
		auto ec = pnm::finalize_to_png(img, outfile.getValue(), format.getValue(),
		                               cmd.getProgramName());
#endif
		params p{};
		if (recreate.isSet()) {
			std::string comment;
			auto ec = extract_comment_from_file(recreate.getValue().c_str(), comment);
			if (ec) {
				std::cerr << "Failed to extract comment from " << recreate.getValue() << ".\n";
				return 1;
			} else {
				std::istringstream cs(comment);
				std::string line;
				cs >> kblib::get_line(line);
				if (not kblib::starts_with(line, "Generated by MapGen v")) {
					std::cerr << "Image comment format does not appear valid\n";
				} 
				// TODO: check version number
				std::string options;
				std::string parameters;
				std::string cmap_name;
				while (cs >> kblib::get_line(line)) {
					if (kblib::starts_with(line, "Options: ")) {
						options = line.substr(9);
					} else if (kblib::starts_with(line, "Parameters: ")) {
						parameters = line.substr(12);
						std::replace(parameters.begin(), parameters.end(), ';', '\n');
					} else if (kblib::starts_with(line, "cmap name: ")) {
						cmap_name = line.substr(11);
					}
				}
				
				std::cout << argv[0] << ' ' << options << "-m " << cmap_name;
			}
		}
		if (paramsfile.isSet()) {
			std::ifstream pf(paramsfile.getValue());
			if (not pf.is_open()) {
				std::cerr << "Could not open "
				          << kblib::quoted(paramsfile.getValue()) << '\n';
				return 1;
			}
			std::string data;
			kblib::get_contents(pf, data);
			std::istrstream is(data.c_str(), data.length());
			// is.exceptions(std::ios::failbit);

			// auto& is = pf;
			try {
				is >> p;
				if (is.bad()) {
					std::cerr << "Error reading params: Bad input\n";
					return 1;
				} else if (is.fail()) {
					std::cerr << "Error reading params\n";
					return 1;
				}
			} catch (std::ios::failure& e) {
				std::cerr << "Error reading parameters: " << e.what() << '\n';
				return 1;
			}
		}
		if (cmapfile.isSet()) {
			current_cmap = read_cmap(cmapfile.getValue().c_str());
		}

		if (list_cmap.getValue() == 1) {
			std::cout << "Number\tName\tSize\n";
			for (const auto i : kblib::range(current_cmap.size())) {
				const auto& [name, pal] = current_cmap[i];
				std::cout << i << ":" << name << '\t' << pal.size() << '\n';
			}
		} else if (list_cmap.getValue()) {
			std::cout << ":Name:Size\nValues\n";
			write_cmap(current_cmap);
		}
		if (convert_srgb.isSet()) {
			if (convert_srgb.getValue() == ":m") {
				cmap_t m_cmaps;
				m_cmaps.push_back({"viridis", sRGB_to_cmap(viridis_srgb_floats)});
				m_cmaps.push_back({"magma", sRGB_to_cmap(magma_srgb_floats)});
				m_cmaps.push_back({"inferno", sRGB_to_cmap(inferno_srgb_floats)});
				m_cmaps.push_back({"plasma", sRGB_to_cmap(plasma_srgb_floats)});
				m_cmaps.push_back({"cividis", sRGB_to_cmap(cividis_srgb_floats)});
				m_cmaps.push_back({"parula", sRGB_to_cmap(parula_srgb_floats)});
				write_cmap(m_cmaps);
			} else {
				auto m = read_cmap_from_file(convert_srgb.getValue());
				write_cmap(m);
			}
		}
		if (auto params_verbosity = list_params.getValue()) {
			if (params_verbosity == 1) {
				std::cout << p;
			} else if (params_verbosity >= 2) {
				p.write_verbose(std::cout);
			}
		}
		if (outfile.isSet()) {
			auto img = genMap<pnm::rgb_pixel>(
			    lookupMap(cmap.getValue()), logarithmic, p, corner.getValue().a(),
			    corner.getValue().b(), width.getValue(), res.getValue());
			std::string comment
			    = kblib::concat("Generated by MapGen v", vnum, ". \nOptions:");
			for (auto arg : gsl::span(argv + 1, argc - 1)) {
				kblib::append(comment, ' ', arg);
			}
			comment += " \nParameters: ";
			{
				std::ostringstream s;
				s << p;
				auto tmp = s.str();
				std::replace(tmp.begin(), tmp.end(), '\n', ';');
				comment += tmp;
			}
			comment += "\ncmap name: ";
			comment += current_cmap[lookupMap(cmap.getValue())].first;
			comment += '\n';
			std::clog << comment << '\n';
			auto ec = pnm::finalize_to_png(img, outfile.getValue(),
			                               format.getValue(), comment);

			if (ec) {
				std::cerr << "Error writing file.\n";
				return ec.value();
			}
		}
		if (save_cmap.isSet()) {
			auto ec = draw_cmaps(save_cmap.getValue(), format.getValue(),
			                     cmapfile.getValue());
			if (ec)
				return ec;
		}
	} catch (TCLAP::ArgException& e) {
		std::cerr << "Argument parsing failed. " << e.what() << '\n';
		return 1;
	}
}