local st = require "util.stanza";

local xmlns_ibb = "http://jabber.org/protocol/ibb";
local xmlns_tcp = "http://prosody.im/protocol/tcpproxy";

local host_attr, port_attr = xmlns_tcp.."\1host", xmlns_tcp.."\1port";

local base64 = require "util.encodings".base64;
local b64, unb64 = base64.encode, base64.decode;

local host = module.host;

local open_connections = {};

local function new_session(jid, sid, conn)
	if not open_connections[jid] then
		open_connections[jid] = {};
	end
	open_connections[jid][sid] = conn;
end
local function close_session(jid, sid)
	if open_connections[jid] then
		open_connections[jid][sid] = nil;
		if next(open_connections[jid]) == nil then
			open_connections[jid] = nil;
		end
		return true;
	end
end

function proxy_component(origin, stanza)
	local ibb_tag = stanza.tags[1];
	if (not (stanza.name == "iq" and stanza.attr.type == "set") 
		and stanza.name ~= "message")
		or
		(not (ibb_tag)
		 or ibb_tag.attr.xmlns ~= xmlns_ibb) then
		if stanza.attr.type ~= "error" then
			origin.send(st.error_reply(stanza, "cancel", "service-unavailable"));
		end
		return;
	end
	
	if ibb_tag.name == "open" then
		-- Starting a new stream
		local to_host, to_port = ibb_tag.attr[host_attr], ibb_tag.attr[port_attr];
		local jid, sid = stanza.attr.from, ibb_tag.attr.sid;
		if not (to_host and to_port) then
			return origin.send(st.error_reply(stanza, "modify", "bad-request", "No host/port specified"));
		elseif not sid or sid == "" then
			return origin.send(st.error_reply(stanza, "modify", "bad-request", "No sid specified"));
		elseif ibb_tag.attr.stanza ~= "message" then
			return origin.send(st.error_reply(stanza, "modify", "bad-request", "Only 'message' stanza transport is supported"));
		end
		local conn, err = socket.tcp();
		if not conn then
			return origin.send(st.error_reply(stanza, "wait", "resource-constraint", err));
		end
		conn:settimeout(0);
		
		local success, err = conn:connect(to_host, to_port);
		if not success and err ~= "timeout" then
			return origin.send(st.error_reply(stanza, "wait", "remote-server-not-found", err));
		end
		
		local listener,seq = {}, 0;
		function listener.onconnect(conn)
			origin.send(st.reply(stanza));
		end
		function listener.onincoming(conn, data)
			origin.send(st.message({to=jid,from=host})
				:tag("data", {xmlns=xmlns_ibb,seq=seq,sid=sid})
				:text(b64(data)));
			seq = seq + 1;
		end
		function listener.ondisconnect(conn, err)
			origin.send(st.message({to=jid,from=host})
				:tag("close", {xmlns=xmlns_ibb,sid=sid}));
			close_session(jid, sid);
		end
		
		conn = server.wrapclient(conn, to_host, to_port, listener, "*a" );
		new_session(jid, sid, conn);
	elseif ibb_tag.name == "data" then
		local conn = open_connections[stanza.attr.from][ibb_tag.attr.sid];
		if conn then
			local data = unb64(ibb_tag:get_text());
			if data then
				conn:write(data);
			else
				return origin.send(
					st.error_reply(stanza, "modify", "bad-request", "Invalid data (base64?)")
				);
			end
		else
			return origin.send(st.error_reply(stanza, "cancel", "item-not-found"));
		end
	elseif ibb_tag.name == "close" then
		if close_session(stanza.attr.from, ibb_tag.attr.sid) then
			origin.send(st.reply(stanza));
		else
			return origin.send(st.error_reply(stanza, "cancel", "item-not-found"));
		end
	end
end

require "core.componentmanager".register_component(host, proxy_component);