#include "../c++/io.hpp"

#include <iostream>

#include "echo.hpp"

saw::error_or<void> handle_echo_write(saw::io_stream<saw::net::Os>& rmt_clt, saw::message& state, uint64_t tbw){
	auto eov = rmt_clt.write(&state.data[state.already_written], tbw);
	if(eov.is_error()){
		return std::move(eov.get_error());
	}

	auto val = eov.get_value();
	state.already_written += val;
	if(state.already_written > state.already_read){
		exit(-1);
	}

	return saw::void_t{};
}

void handle_echo_message(saw::io_stream<saw::net::Os>& rmt_clt, bool& keep_running, saw::message& state){
	rmt_clt.read_ready().then([&](){
		for(;;){
			uint64_t tbr = state.data.size() < state.already_read ? 0: state.data.size() - state.already_read;
			if(tbr == 0){
				exit(-1);
			}
			auto eov = rmt_clt.read(&state.data[state.already_read], tbr);

			if(eov.is_error()){
				auto& err = eov.get_error();
				if(err.is_critical()){
					std::cerr<<err.get_category()<<std::endl;
					exit(err.get_id());
				}else{
					break;
				}
			}
			auto& read_bytes = eov.get_value();

			std::cout<<"Read "<<read_bytes<<" bytes"<<std::endl;
			
			if(read_bytes == 0u){
				exit(-1);
			}

			bool trigger_write = (state.already_read == state.already_written);

			state.already_read += read_bytes;
			if(state.already_read > state.data.size()){
				state.already_read = state.data.size();
			}
			if(trigger_write){
				auto eov = handle_echo_write(rmt_clt, state, state.already_read - state.already_written);
				if(eov.is_error()){
					auto& err = eov.get_error();
					if(err.is_critical()){
						std::cerr<<"After triggered write: "<<err.get_category()<<std::endl;
						exit(err.get_id());
					}else {
						break;
					}
				}
			}
		}
	}).detach();

	rmt_clt.write_ready().then([&](){
		for(;;){
			if(state.already_read < state.already_written){
				exit(-1);
			}
			uint64_t tbw = state.already_read - state.already_written;
			if(tbw == 0){
				break;
			}

			auto eov = handle_echo_write(rmt_clt, state, tbw);
			if(eov.is_error()){
				auto& err = eov.get_error();
				if(err.is_critical()){
					std::cerr<<err.get_category()<<std::endl;
					exit(err.get_id());
				}else {
					break;
				}
			}
		}
	}).detach();

	rmt_clt.on_read_disconnected().then([&](){
		keep_running = false;
	}).detach();
}

int main(){
	auto eo_aio = saw::setup_async_io();
	if(eo_aio.is_error()){
		auto& err = eo_aio.get_error();
		std::cerr<<err.get_message()<<std::endl;
		return err.get_id();
	}
	auto& aio = eo_aio.get_value();
	/**
	 * Make the event loop the current event loop on this thread
	 */
	saw::wait_scope wait_scope{aio.event_loop};

	bool keep_running = true;
	aio.event_port.on_signal(saw::Signal::Terminate).then([&keep_running](){
		keep_running = false;
	}).detach();

	auto& network = aio.io->get_network();
	saw::own<saw::network_address<saw::net::Os>> addr = nullptr;
	saw::own<saw::server<saw::net::Os>> srv = nullptr;
	saw::own<saw::io_stream<saw::net::Os>> remote_client = nullptr;

	saw::message msg_state;

	std::cout<<"Starting to resolve address"<<std::endl;
	/**
	 * Try to resolve address. If resolved
	 */
	network.resolve_address(saw::echo_address, saw::echo_port).then([&](auto net_addr){
		std::cout<<"Resolved address"<<std::endl;
		addr = std::move(net_addr);
		srv = network.listen(*addr);
		if(srv){
			srv->accept().then([&](auto client) -> saw::error_or<void>{

				if(!remote_client){
					std::cout<<"Accepted client"<<std::endl;
					remote_client = std::move(client);

					if(remote_client){
						std::cout<<"Spinning up handler"<<std::endl;
						handle_echo_message(*remote_client, keep_running, msg_state);
					}else{
						keep_running = false;
					}
				}
				return saw::make_error<saw::err::critical>();
			}).detach();
		}else{
			keep_running = false;
		}
	}).detach([&](auto err) {
		std::cout<<"Failed to resolve address: "<<err.get_category()<<":"<<err.get_message()<<std::endl;
		keep_running = false;
		return err;
	});

	std::cout<<"Entering waiting loop"<<std::endl;

	wait_scope.poll();
	while(keep_running){
		wait_scope.wait_for(1*1000*1000);
	}

	std::cout<<"\n\nShutting down echo server"<<std::endl;

	return 0;
}