module bencode;

import std.variant: Algebraic;
import std.conv: to;
import std.algorithm.searching: countUntil;
import std.array: appender, replicate;
import std.meta: AliasSeq, staticIndexOf;
import std.bigint: BigInt;

private class BEncodeException : Exception
{
    this(string msg, string file = __FILE__, size_t line = __LINE__) pure nothrow @safe @nogc
    {
        super(msg, file, line);
    }
}

private BStr getBStr(ref ubyte[] file) pure @trusted
{
	auto i = countUntil(file, ':');
	
	if (i == -1)
		throw new BEncodeException("No : symbol");
		
	auto len = to!int(cast(string)file[0 .. i]);
	
	if (len+i+1 > file.length)
		throw new BEncodeException("String too long");
	
	auto ret = file[i+1 .. len+i+1];
	file = file[len+i+1 .. $];
	return cast(string)ret;
}

private BInt getInt(ref ubyte[] file) pure @trusted
{
	if (file[0] != 'i')
		throw new BEncodeException("No i symbol");
		
	auto i = countUntil(file, 'e');
	
	if (i == -1)
		throw new BEncodeException("No e symbol");
	
	BigInt ret;
	try {
		ret = BigInt(cast(string)file[1 .. i]);
	} catch (Exception ex) {
		throw new BEncodeException(ex.msg);
	}
	
	file = file[i+1 .. $];
	return ret;
}

private enum EL_TYPE
{
	BYTE_STR,
	INT,
	LIST,
	DICT,
	END
}

private EL_TYPE GetType(ref ubyte[] file) pure @safe
{
	switch (file[0]) with(EL_TYPE) {
		case '0': .. case '9':
			return BYTE_STR;
		case 'i':
			return INT;
		case 'l':
			return LIST;
		case 'd':
			return DICT;
		case 'e':
			return END;
		default:
			throw new BEncodeException("Unknown element type");
	}
}

public struct OrderedAA(K, V) {
	import std.array;
	
	public struct KV
	{
		private K _key;
		private V _value;
		
		@property K key() pure nothrow @safe @nogc
		{
			return _key;
		}
		
		@property void key(K val) pure nothrow @safe @nogc
		{
			_key = val;
		}
		
		@property ref V value() pure nothrow @safe @nogc
		{
			return _value;
		}
		
		@property void value(V val) @trusted
		{
			_value = val;
		}
		
		this(K k, V v) pure nothrow @safe @nogc
		{
			_key = k;
			_value = v;
		}
	}
	
	private KV[] list;
	
	KV[] byKeyValue() pure nothrow @nogc
	{
	    return list;
	}
	
	auto byKey() pure nothrow @nogc
	{
		import std.algorithm.iteration;
	    return list.map!(x => x.key);
	}
	
	auto byValue() pure nothrow @nogc
	{
	    import std.algorithm.iteration;
	    return list.map!(x => x.value);
	}
	
	void opIndexAssign(V v, K k) nothrow @safe
	{
		auto i = countUntil!"a.key == b"(list, k);
		if (i == -1) {
			list ~= KV(k, v);
		} else {
			try {
				list[i].value = v;
			} catch (Exception ex) {
				assert(false, ex.msg);
			}
		}
	}
	
	KV *opBinaryRight(string op : "in")(K k) @trusted
	{
		auto i = countUntil!"a.key == b"(list, k);
		if (i == -1)
			return null;
		return &list[i];
	}
	
	KV opIndex(K k) @trusted
	{
		auto pKV = opBinaryRight!"in"(k);
		
		if (pKV is null)
			throw new BEncodeException("Item with key \"" ~ cast(string)k ~ "\" not found");
		else return *pKV;
	}
}

alias BInt = BigInt;
alias BStr = string;
alias BDict = OrderedAA!(BStr, BElement);
alias BList = BElement[];

alias BContent = Algebraic!(BAllowedTypes);

alias BAllowedTypes = AliasSeq!(BStr, BInt, BDict, BList, BElement);

private enum isInAllowed(T)() {
	return staticIndexOf!(T, BAllowedTypes) != -1 || isNumeric!T || is(T == BContent);
}

import std.stdio;
import std.traits;
import std.range.primitives;

private BElement parseRecursive(T)(T inp)
{
	auto base = new BElement();
	base.elSet(parseIRecursive(inp));
	return base;
}

private BContent parseIRecursive(T)(T inp, int printLevel = 0)
{
	static if (isInAllowed!T) {
		static if (isNumeric!T)
			return BContent(BigInt(to!string(inp)));
		else
			return BContent(inp);
		
		
	} else static if (isAssociativeArray!T) {
		BDict dict;
		
		foreach (item; inp.byKeyValue()) {
			BElement newEl = new BElement();
			newEl.elSet(parseIRecursive(item.value, printLevel + 1));
			dict[item.key] = newEl;
		}
		
		return BContent(dict);
	} else static if (isInputRange!T && !isSomeString!T) {
		BList list;
		
		foreach (item; inp) {
			BElement newEl = new BElement();
			newEl.elSet(parseIRecursive(item, printLevel + 1));
			list ~= newEl;
		}
		
		return BContent(list);
	} else
		static assert(false, "Unsupported type " ~ fullyQualifiedName!T);
}

public class BElement
{
	private BContent el;
	
	public void elSet(T)(T set) {
		if (!isInAllowed!T)
			throw new BEncodeException("Cannot store " ~ fullyQualifiedName!T);
		el = set;
	}
	
	public this() {  }
	public this(V)(V content) if (!isInAllowed!V) { el = parseRecursive(content); }
	public this(V)(V content) if (isInAllowed!V) {
		static if (isNumeric!V)
			el = BigInt(to!string(content));
		else
			el = content;
	}
	
	private string toIString(BElement el, int level)
	{
		auto ret = appender!string();
		
		if (el.el.type == typeid(BStr)) {
			ret.put("\t".replicate(level) ~ cast(string)*el.str() ~ "\n");
		} else if (el.el.type == typeid(BInt)) {
			ret.put("\t".replicate(level) ~ to!string(*el.integer()) ~ "\n");
		} else if (el.el.type == typeid(BDict)) {
			foreach (item; (*el.dict()).byKeyValue()) {
				ret.put("\t".replicate(level) ~ cast(string)item.key ~ ": " ~ "\n");
				ret.put(toIString(item.value, level + 1));
			}
		} else if (el.el.type == typeid(BList)) {
			foreach (item; *el.list()) {
				ret.put(toIString(item, level + 1));
			}
		} else if (el.el.type == typeid(BElement)) {
			ret.put(toIString(*el.element(), level));
		}
	
		return ret.data;
	}
	
	override string toString() @trusted
	{
		return toIString(this, 0);
	}
	
	public BStr serialize()
	{
		BStr serializeBStr(BStr inp) {
			return to!string(inp.length) ~ ':' ~ inp;
		}
		
		auto ret = appender!BStr();
		
		if (el.type == typeid(BStr)) {
			ret.put(serializeBStr(*bstr()));
		} else if (el.type == typeid(BInt))
			ret.put('i' ~ to!string(*integer()) ~ 'e');
		else if (el.type == typeid(BDict)) {
			ret.put('d');
			foreach (item; (*dict()).byKeyValue()) {
				ret.put(serializeBStr(item.key));
				ret.put(item.value.serialize());
			}
			ret.put('e');
		} else if (el.type == typeid(BList)) {
			ret.put('l');
			foreach (item; *list()) {
				ret.put(item.serialize());
			}
			ret.put('e');
		} else if (el.type == typeid(BElement)) {
			ret.put((*element()).serialize());
		}
	
		return ret.data;
	}
	
	public BElement opIndex(BStr v) {
		if (el.type != typeid(BDict))
			throw new BEncodeException("Element is not a dictionary");
		
		BDict.KV *kv = (*dict()).opBinaryRight!"in"(v);
		
		if (kv is null)
			return null;
		else
			return (*kv).value;
	}
	
	public BElement opIndex(uint v) {
		assert(el.type == typeid(BDict) || el.type == typeid(BList), "Element is not a dictionary");
		
		if (el.type == typeid(BDict))
			return ((*dict()).byValue())[v];
		else
			return (*list())[v];
	}
	
	BStr *bstr() { return el.peek!BStr(); }
	string *str() { return cast(string*)bstr(); }
	BInt *integer() { return el.peek!BInt(); }
	BDict *dict() { return el.peek!BDict(); }
	BList *list() { return el.peek!BList(); }
	BElement *element() { return el.peek!BElement(); }
	
	void opAssign(BStr str) { el = str; }
	void opAssign(BInt integer) { el = integer; }
	void opAssign(BDict dict) { el = dict; }
	void opAssign(BList list) { el = list; }
	
	void opOpAssign(string op : "~", V)(V v)
	{
		assert(el.type == typeid(BList), "Base type is not a list");
		
		*list() ~= parseRecursive(v);
	}
	
	private auto innerType(T)(T val) {
		static if (is (T == BElement)) {
			if (val.el.type == BElement)
				return innerType(val.el.get!BElement());
		}
		return val.type;
	}
	
	void opIndexOpAssign(string op : "~", V, K)(V v, K k)
	{
		assert(el.type == typeid(BDict), "Base type is not a dictionary");
		assert(is(K == typeof((*dict()).byKey()[0])), "Dictionary keys are not of type " ~ fullyQualifiedName!K ~ " (" ~ fullyQualifiedName!(typeof((*dict()).byKey()[0])) ~ " expected)");
		
		auto i = k in *dict();
		
		if (i !is null) {
			if (innerType((*i).value.el) == typeid(BList)) {
				(*i).value.opOpAssign!("~", V)(v);
			} else {
				static if (isAssociativeArray!V) {
					auto aa = (*i).value.dict();
					foreach(item; v.byKeyValue()) {
						(*aa)[item.key] = parseRecursive(item.value);
					}
				} else
					throw new BEncodeException("Cannot add non-associative array to dictionary");
			}
		} else
			(*dict())[k] = parseRecursive(v);
	}
}

private BElement _iParse(T)(ref ubyte[] file, ref uint sC, ref uint eC)
{
	BStr curKey = null;
	BElement ret = new BElement();
	ret.el = T.init;

	static if (is(T == BStr)) {
		ret.el = getBStr(file);
		
		if (file.length != 0)
			throw new BEncodeException("Unexpected data after a string");
		
		return ret;
	} else static if (is(T == BInt)) {
		ret.el = getInt(file);
		
		if (file.length != 0)
			throw new BEncodeException("Unexpected data after a int");
		
		return ret;
	}

	while (file.length != 0) {
		auto type = GetType(file);
		
		BElement el;
		
		switch (type) with (EL_TYPE) {
			case BYTE_STR:
			case INT:
				if (type == BYTE_STR)
					el = new BElement(getBStr(file));
				else
					el = new BElement(getInt(file));
			
				if (curKey is null) {
					if (ret.el.type == typeid(BList)) {
						(*ret.list()) ~= el;
						continue;
					}
					
					if (type == BYTE_STR)
						curKey = *el.str();
					else
						throw new BEncodeException("Integer key at " ~ cast(string)file[0 .. $] ~ "...");
						
					continue;
				}
				
				if (ret.el.type == typeid(BDict))
					(*(ret.dict()))[curKey] = el;
				else
					throw new BEncodeException("Invalid element base container");
					
				curKey = null;
				break;
			case LIST:
			case DICT:
				sC++;
				file = file[1 .. $];
			
				if (type == LIST)
					el = new BElement(_iParse!BList(file, sC, eC).el);
				else
					el = new BElement(_iParse!BDict(file, sC, eC).el);
				
				if (ret.el.type == typeid(BList))
					ret.el ~= el;
				else if (ret.el.type == typeid(BDict)) {
					if (curKey is null)
						throw new BEncodeException("No key on dictionary");
					else
						(*(ret.dict()))[curKey] = el;
						
					curKey = null;
				} else
					throw new BEncodeException("Invalid element base container");

				curKey = null;
				break;
			case END:
				eC++;
				file = file[1 .. $];
				return ret;
			default:
				throw new BEncodeException("Unknown element type");
		}
	}

	return ret;
}

BElement bencodeParse(ubyte[] file)
{
	auto type = GetType(file);
	
	uint startCount = 0, endCount = 0;
	
	BElement ret;
	
	final switch (type) with (EL_TYPE) {
		case BYTE_STR:
			ret = _iParse!BStr(file, startCount, endCount);
			break;
		case INT:
			ret = _iParse!BInt(file, startCount, endCount);
			break;
		case LIST:
			startCount++;
			file = file[1 .. $];
			ret = _iParse!BList(file, startCount, endCount);
			break;
		case DICT:
			startCount++;
			file = file[1 .. $];
			ret = _iParse!BDict(file, startCount, endCount);
			break;
		case END:
			throw new BEncodeException("End of unstarted list/dict");
	}
	
	if (startCount != endCount)
		throw new BEncodeException("d/l and e count doesn't match (start " ~ to!string(startCount) ~ ", end " ~ to!string(endCount) ~ ")");
		
	return ret;
}

unittest
{
	import std.file: read;
	import std.algorithm.comparison: equal;
	import std.exception: assertThrown;
	
	auto file = cast(ubyte[])read("OnlyHuman-e-0.21.7-20170411.torrent");
	
	void testNumber() {
		static auto c = 0;
		c++;
		writeln("test ", c);
	}
	
	testNumber();
	auto bencode = bencodeParse(file);
	
	testNumber();
	assert(equal(bencode.serialize(), file));
	
	testNumber();
	assert(equal(*bencode["comment"].str(), "OnlyHuman-e-0.21.7-20170411"));
	
	testNumber();
	assert(equal(*bencode["info"]["files"][6]["path"][0].str(), "OnlyHuman-e-0.21.7-20170411.txt"));
	
	testNumber();
	// Trying to add value without key
	assertThrown!BEncodeException(bencode["info"] ~= "test");
	
	testNumber();
	// Add item to list - OK
	bencode["info"]["files"] ~= "test";
	
	testNumber();
	// Add list to list - OK
	bencode["info"]["files"] ~= [ "test2", "test3" ];
	
	testNumber();
	// Add dictionary to list - OK
	bencode["info"]["files"] ~= [ "test4": new BElement(1000), "test5": new BElement([ "very nice" ]) ];
	
	testNumber();
	assert(equal(bencode.serialize(), cast(string)read("mod1.torrent")));
	
	bencode = new BElement(BDict.init);
	
	testNumber(); // 10
	// Add key (length) and value (1000) to root dictionary
	bencode["length"] ~= 1000;
	
	testNumber();
	// Add key (very) and value (nicep) to "path" dictionary
	bencode["path"] ~= [ "very": "nicep" ];
	
	testNumber();
	// Add key (very good) and value (nice) to "path" dictionary
	bencode["path"] ~= [ "very good": "nice" ];
	
	testNumber();
	// Trying to add value to "path" dictionary
	assertThrown!BEncodeException(bencode["path"] ~= "nice");
	
	testNumber();
	// Overwrite nicep
	bencode["path"] ~= [ "very": "nice" ];
	
	testNumber();
	assert(bencode.serialize() == "d6:lengthi1000e4:pathd4:very4:nice9:very good4:niceee");
	
	testNumber();
	assert(bencodeParse(cast(ubyte[])"d4:spam3:egge").serialize() == "d4:spam3:egge");
	
	testNumber(); // 17
	assert(bencodeParse(cast(ubyte[])"de").serialize() == "de");
	
	auto invalid = [ "d", "ade", ":de" /* 20 */, "-de", "1de", "d4:spam3:egg", "d ", "l-something-e", "4", "a", ":", "-", "." /* 30 */, "e", "l", "l4:spam", "l ", "l:",
		"1e", "1234567890e", "iasdfe", "ie", "5:spam" /* 40 */, "100000:spam", "4spam", "10spam", "4;spam", "spam", "interesting", "0", "1:spam" ];
	auto valid = [ "l4:spame", "l4:spami42ee" /* 50 */, "le", "i0e", "i1e", "i3e", "i42e", "i100e", "i1234567890e", "i165183413841846846518341685135165e",
		"i-0e", "i-1e" /* 60 */, "i-3e", "i-42e", "i-100e", "i-1234567890e", "i-165183413841846846518341685135165e", "4:spam", "9:spam:eggs", "3:\x00\x01\x00", "i-e" /* 69 */, "d0:i1ee" ];
	
	foreach (item; invalid) {
		testNumber();
		assertThrown!BEncodeException(bencodeParse(cast(ubyte[])item));
	}
		
	foreach (item; valid) {
		testNumber();
		bencodeParse(cast(ubyte[])item);
	}
}