/* c++17 */

#include <vector>
#include <unordered_map>
#include <string>
#include <string_view>
#include <system_error>

#include <cstdint>
#include <cassert>
#include <cstdio>

#include <err.h>

#include <afp/finder_info.h>


#include "omf.h"
#include "rel.h"

void save_omf(const std::string &path, std::vector<omf::segment> &segments, bool compress, bool expressload);

/* since span isn't standard yet */
typedef std::basic_string_view<uint8_t> byte_view;


struct symbol {
	std::string name;
	std::string file;
	uint32_t value = 0;
	unsigned id = 0;
	unsigned count = 0;

	bool absolute = false;
	bool defined = false;
};

struct pending_reloc : public omf::reloc {
	unsigned id = 0;
};

std::unordered_map<std::string, unsigned> symbol_map;
std::vector<symbol> symbol_table;


std::vector<omf::segment> segments;

/* nb - pointer may be invalidated by next call */
symbol *find_symbol(const std::string &name) {
	
	auto iter = symbol_map.find(name);
	if (iter != symbol_map.end()) return &symbol_table[iter->second - 1];

	unsigned id = symbol_table.size() + 1;
	symbol_map.emplace(name, id);

	auto &rv = symbol_table.emplace_back();
	rv.name = name;
	rv.id = id;
	return &rv;
}



struct cookie {
	std::string file;
	std::vector<unsigned> remap;
	uint32_t start = 0;
	uint32_t length = 0;
};

void process_labels(byte_view &data, cookie &cookie) {

	for(;;) {
		assert(data.size());
		unsigned flag = data[0];
		if (flag == 0x00) return;

		unsigned length = flag & 0x1f;
		assert(length != 0);
		assert(data.size() >= length + 4);

		std::string name(data.data() + 1, data.data() + 1 + length);
		data.remove_prefix(1 + length);
		uint32_t value = data[0] | (data[1] << 8) | (data[2] << 16);
		data.remove_prefix(3);

		symbol *e = find_symbol(name);
		switch (flag & ~0x1f) {
			case SYMBOL_EXTERNAL:
				/* map the unit symbol # to a global symbol # */
				value &= 0x7fff;
				if (cookie.remap.size() < value + 1)
					cookie.remap.resize(value + 1);
				cookie.remap[value] = e->id;
				break;


			case SYMBOL_ENTRY+SYMBOL_ABSOLUTE:
				if (e->defined && e->absolute && e->value == value)
					break; /* allow redef */

			case SYMBOL_ENTRY:
				if (e->defined) {
					warnx("%s previously defined (%s)", e->name.c_str(), e->file.c_str());
					break;
				}
				e->defined = true;
				e->file = cookie.file;
				if (flag & SYMBOL_ABSOLUTE) {
					e->absolute = true;
					e->value = value;
				} else {
					e->absolute = false;
					e->value = value - 0x8000 + cookie.start;
				}
				break;
			default:
				errx(1, "%s: Unsupported flag: %02x\n", cookie.file.c_str(), flag);
				break;
		}
	}
}


void process_reloc(byte_view &data, cookie &cookie) {

	auto &seg_data = segments.back().data;

	for(;;) {
		assert(data.size());
		unsigned flag = data[0];
		if (flag == 0x00) return;

		assert(data.size() >= 4);

		uint32_t offset = data[1] | (data[2] << 8);
		unsigned x = data[3];
		data.remove_prefix(4);

		offset += cookie.start;
		bool external = false;
		unsigned shift = 0;
		uint32_t value = 0;
		unsigned size = 0;

		if (flag == 0xff) {
			/* shift */
			assert(data.size() >= 4);
			unsigned flag = data[0];
			value = data[1] | (data[2] << 8) | (data[3] << 16);
			value -= 0x8000;
			external = flag & 0x04;
			switch(flag & ~0x04) {
				case 0xd0:
					shift = 16;
					size = 1;
					break;
				case 0xd1:
					shift = 8;
					size = 2;
					break;
				case 0xd3:
					shift = 8;
					size = 1;
					break;
				default: /* bad */
					errx(1, "%s: Unsupported flag: %02x\n", cookie.file.c_str(), flag);
					break;
			}
		} else {
			assert((flag  & ~(0x0f|0x10|0x20|0x80)) == 0);

			unsigned size = 0;
			switch(flag & (0x80 + 0x20)) {
				case 0:
					size = 1;
					assert(offset + 0 < cookie.length);
					break;
				case 0x20:
					size = 3;
					assert(offset + 2 < cookie.length);
					break;
				case 0x80:
					size = 2;
					assert(offset + 1 < cookie.length);
					break;
				default: /* bad size */
					errx(1, "%s: Unsupported flag: %02x\n", cookie.file.c_str(), flag);
					break;
			}
			external = flag & 0x10;

			switch(size) {
				case 3: value |= seg_data[offset+2] << 16;
				case 2: value |= seg_data[offset+1] << 8;
				case 1: value |= seg_data[offset+0];
			}


			if (size > 1) value -= 0x8000;
			value += cookie.start;

		}

		/* external resolutions are deferred for later */

		if (external) {
			/* x = local symbol # */
			pending_reloc r;
			assert(x < cookie.remap.size());
			r.id = cookie.remap[x]; /* label reference is 0-based */
			r.size = size;
			r.offset = offset;
			r.value = value;
			r.shift = shift;

			relocations.emplace_back(r);
		} else {
			uint32_t value = 0;
			omf::reloc r;
			r.size = size;
			r.offset = cookie.start; /* ???? */
			r.value = value;
			r.shift = shift;

			seg.relocs.emplace_back(r);
		}

	}
}

/*
void add_libraries() {
	auto iter = libs.begin();
	auto end = libs.end();

	for(;;) {



	}
}
*/


void process_unit(const std::string &path) {

	cookie cookie;
	/* skip over relocs, do symbols first */


	std::error_code ec;
	mapped_file mf(path, mapped_file::readonly, ec);
	if (ec) {
		errx(1, "Unable to open %s: %s" path.c_str(), ec.message().c_str());
	}


	afp::finder_info fi;

	fi.read(path, ec);

	if (ec) {
		errx(1, "Error reading filetype %s: %s", path.c_str(), ec.message().c_str());
	}

	if (fi.prodos_file_type() != 0xf8) {
		errx(1, "Wrong file type: %s", path.c_str());
	}

	uint32_t offset = fi.prodos_aux_type();
	if (offset+2 > mf.size()) {
		errx(1, "Invalid aux type %s", path.c_str());
	}


	omf::segment &seg = segments.back();

	seg.data.append(mf.data(), mf.data() + offset);

	byte_view data(mf.data() + offset, mf.size() - offset);

	cookie.start = seg.data.size();
	cookie.length = offset;
	cookie.file = path;

	byte_view rr = data;
	/* skip over the relocation records so we can process the labels first. */
	/* this is so external references can use the global symbol id */
	for(;;) {
		assert(data.size() >= 6);
		if (data[0] == 0) break;
		data.remove_prefix(4);
	}
	data.remove_prefix(1);
	process_labels(data, cookie);
	assert(data.size() == 1);

	/* now relocations */
	process_reloc(rr, cookie);
}


void finalize(void) {

	for (auto &r in relocations) {
		assert(r.id <= symbol_map.length());
		const auto &label = symbol_map[rr.id];

		r.value = label.value;
		seg.relocs.emplace_back(r);
	}
	relocations.clear();
}

void print_symbols(void) {

	if (symbol_table.empty()) return;

	/* alpha */
	std::sort(symbol_table.begin(), symbol_table.end(),
		[](const symbol &a, const symbol &b){
			return a.name < b.name; 
		});

	for (const auto &lab : symbol_table) {
		fprintf(stdout, "%-20s: $%06x\n", lab.name.c_str(), lab.value);
	}
	fputs("\n", stdout);
	/* numeric */
	std::sort(symbol_table.begin(), symbol_table.end(),
		[](const symbol &a, const symbol &b){
			return a.value < b.value; 
		});

	for (const auto &lab : symbol_table) {
		fprintf(stdout, "%-20s: $%06x\n", lab.name.c_str(), lab.value);
	}
}