use core:io;
use core:net;

/**
 * A basic HTTP server.
 *
 * Listens for incoming HTTP requests on a given socket and calls the `onRequest` member to handle
 * the request.
 */
class Server {
	// Currently active listeners.
	private WeakSet<Listener> listeners;

	// All clients that are currently connected.
	private WeakSet<Client> clients;

	// Regex that matches the end of the header.
	private lang:bnf:Regex endOfHeader = lang:bnf:Regex(".*\r\n\r\n");

	// Timeout to apply to new incoming connections.
	private Duration timeoutDuration = 60 s;

	// Cached string to reply to keepalive messages.
	private Str keepAliveString;

	// Maximum size of a request (1 MB by default).
	Nat maxRequest = 1024n * 1024n;

	// Buffer size for clients.
	Nat clientBuffer = 4096;

	// Create.
	init() {
		init {}
	}

	// Set the timeout.
	assign timeout(Duration d) {
		this.timeoutDuration = d;
		this.keepAliveString = "timeout=${d.inS}; max=10";
	}

	// Run the server on `port` until `close` is called.
	void run(Nat port) {
		unless (listener = listen(port, true)) {
			throw HttpError("Failed to listen on port ${port}.");
		}

		listeners.put(listener);

		while (socket = listener.accept()) {
			Client c(this, socket);
			spawn c.main();
		}

		listeners.remove(listener);

		// Wait for any lingering connections to close.
		Bool anyRunning = true;
		while (anyRunning) {
			yield();
			anyRunning = false;
			for (c in clients)
				anyRunning |= c.running;
		}
	}

	// Close the listener inside the server. This also causes 'run' to return if another thread is
	// blocked inside it.
	void close() {
		// Close all listeners.
		for (l in listeners)
			l.close();

		// Close all clients as well.
		for (c in clients)
			c.close();
	}

	// Called when a request was received. Expected to return an appropriate response.
	Response onRequest(Request request) : abstract;

	// Override to handle errors from clients.
	Response onServerError(Exception e) : abstract;

	// Called when an immediate response should be generated. Allows generating an error page.
	Response onError(Status code) : abstract;

	// Called before sending a request, allows a top-level hook for handling headers etc.
	// Returns 'true' if the connection should be kept open. False otherwise.
	Bool prepareResponse(Request request, Response response) {
		// Use the version of the request by default.
		response.version = request.version;

		// Handle keep alive.
		Bool keepAlive = false;
		if (value = request.header("connection"))
			keepAlive = noCaseEqual(value, "keep-alive");

		if (keepAlive) {
			response.header("connection", "keep-alive");
			response.header("keep-alive", keepAliveString);
		}

		return keepAlive;
	}


	/**
	 * State associated with each client.
	 */
	private class Client {
		// The server, so that we can inform about errors and responses.
		Server server;

		// The socket we are using.
		NetStream socket;

		// Input stream.
		NetIStream is;

		// Output stream.
		BufferedOStream os;

		// Buffer from the input stream. To make it possible to handle connections where the client
		// sends multiple requests at once.
		Buffer inputBuffer;

		// Is the main thread running?
		Bool running;

		// Create.
		init(Server server, NetStream socket) {
			init {
				server = server;
				socket = socket;
				is = socket.input;
				os = BufferedOStream(socket.output, server.clientBuffer);
				inputBuffer = buffer(server.clientBuffer);
			}

			is.timeout = server.timeoutDuration;
			server.clients.put(this);
		}

		// Run the client handling.
		void main() {
			running = true;
			try {
				mainBody();
			} catch (Exception e) {
				server.onServerError(e).write(os);
				os.flush();
			}

			server.clients.remove(this);
			running = false;
			os.flush();
			close();
		}

		// Close the client, aborts processing.
		void close() {
			socket.close();
		}

		// Body of the main function, protected by try/catch.
		private void mainBody() {
			Bool exit = false;
			while (is.more & is.error.empty & os.error.empty & !exit) {
				var request = readRequest();
				unless (request)
					return;

				var response = if (request.immediateResponse != Status:none) {
					// Send a response, but close connection afterwards.
					exit = true;
					server.onError(request.immediateResponse);
				} else {
					server.onRequest(request);
				};

				exit |= !server.prepareResponse(request, response);
				exit |= response.write(os);
				os.flush();
			}
		}

		// Read a request from the socket.
		private Request? readRequest() {
			Request? r = readHeader();
			unless (r)
				return null;
			if (r.immediateResponse != Status:none)
				return r;

			// See if there is a body.
			if (bodyLengthStr = r.header("content-length")) {
				if (bodyLength = bodyLengthStr.nat) {
					readBody(bodyLength, r);
				}
			}

			return r;
		}

		private Request? readHeader() {
			while (is.more) {
				// See if we have \r\n\r\n that ends the header. We do this in two steps since the
				// regex matching is cheaper than running the full parser on the buffer first.
				if (server.endOfHeader.match(inputBuffer)) {
					// We have the entire header at least.
					var result = parseRequestHeader(inputBuffer);
					unless (request = result.value) {
						print(inputBuffer.fromUtf8);
						print(result.error.toS);
						return Request(Status:Bad_Request);
					}

					// Let the parser determine the actual end of the request.
					inputBuffer.shift(result.end);
					return request;
				}

				Status error = readData();
				if (error != Status:none) {
					if (inputBuffer.filled == 0)
						return null;
					else
						return Request(error);
				}
			}

			return Request(Status:Bad_Request);
		}

		private void readBody(Nat count, Request to) {
			while (is.more) {
				if (inputBuffer.filled == count) {
					to.body = inputBuffer;
					inputBuffer = Buffer();
					return;
				} else {
					to.body = inputBuffer.cut(0, count);
					inputBuffer.shift(count);
					return;
				}

				Status error = readData();
				if (error != Status:none) {
					to.immediateResponse = error;
					return;
				}
			}

			to.immediateResponse = Status:Bad_Request;
		}

		// Read some data from the socket to the buffer.
		private Status readData() {
			if (inputBuffer.free == 0)
				if (!grow)
					return Status:Request_Entity_Too_Large;

			var oldFilled = inputBuffer.filled;
			inputBuffer = is.read(inputBuffer);
			if (inputBuffer.filled == oldFilled)
				return Status:Request_Timeout;

			return Status:none;
		}

		// Grow the input buffer.
		private Bool grow() {
			if (inputBuffer.count >= server.maxRequest)
				return false;

			Nat newCount = min(inputBuffer.count + server.clientBuffer, server.maxRequest);
			inputBuffer = inputBuffer.grow(newCount);
			true;
		}
	}

}