#!/bin/zsh

# Exit on error e.g. unknown command
set -e
#set -x

# Load ztcp, for opening TCP ports
zmodload zsh/net/tcp
# Load zselect, for testing which file descriptors have activity
zmodload zsh/zselect
# Load zparseopts
zmodload zsh/zutil

# Set server parameters
PORT=8080
HOST=127.0.0.1
HTROOT="$(pwd)"
# TODO: if no -p is specified and 8080 is taken, then try again on a random port

CR=$'\r'
LF=$'\n'
CRLF=$'\r\n'
SERVER="zhttpd/0.0.0 (<https://github.com/awwright/zhttpd>) zsh/$(zsh --version|cut -c 5-)"
MIME_TYPES_PATH="/etc/apache2/mime.types"
typeset -A ext_mediatype
DEFAULT_MEDIA_TYPE="application/octet-stream"

main() {
	if [[ -f "$MIME_TYPES_PATH" ]]; then
		while IFS=$'\t ' read -r mt extensions; do
			# Skip blank lines and comments
			[[ -z $mt || $mt == '#'* ]] && continue
			# Split extensions and map each to the MIME type
			for ext in ${(s: :)extensions}; do
				#echo "${ext:l} -> $mt"
				ext_mediatype[${ext:l}]=$mt
			done
		done < "$MIME_TYPES_PATH"
	else
		echo "Error: mime.types file not found" >&2
	fi

	# Create a TCP listening socket
	# This command stores the fd of the socket in $REPLY
	# This should be closed when finished
	ztcp -l $PORT
	local listen_fd=$REPLY

	debug main "listen_fd=$listen_fd"
	if (( $listen_fd == 0 )); then
		print "Failed to bind to port $PORT"
		exit 1
	fi

	print "Server listening on $HOST:$PORT..."
	print "http://$HOST:$PORT/"

	# Main event loop
	# TODO: use zselect
	# <https://zsh.sourceforge.io/Doc/Release/Zsh-Modules.html#The-zsh_002fzselect-Module>
	while true; do
		#echo zselect -r $listen_fd ${(k)listen_fd_child} -A ready_fds
		zselect -r $listen_fd ${(k)listen_fd_child} -A ready_fds
		for fd in ${(k)ready_fds}; do
			if [[ $fd == $listen_fd ]]; then
				# This will kick off a subprocess in a way that will allow reading and writing to it directly
				coproc handle_client $client_fd
				# Move the fd known as p to a new file descriptor so we can read messages from the coprocess
				exec {fd}<& p
				listen_fd_child[$fd]="$fd"
			else
				# This must be a message from a child process to be handled
				# If read returns nonzero, this means EOF was reached, so close the file descriptor
				# `exec {fd}>&-` closes the fd from the coprocess, so that the number can be reused
				read out <&$fd || { unset "listen_fd_child[$fd]"; exec {fd}>&- }
				echo fd=$fd out=$out
			fi
		done
	done

	# Close the listening socket opened by `ztcp -l $PORT`
	# Otherwise it may be difficult to start this process back up listening on the same port
	ztcp -c $listen_fd
}

# New function to handle a single client connection, including persistent requests
handle_client() {
	# Wait for a connection on $listen_fd
	# This command stores the fd of the client connection socket in $REPLY
	ztcp -a $listen_fd
	local client_fd=$REPLY

	# request-line = method SP request-target SP HTTP-version
	local request_method=
	read -rsd ' ' request_method <&$client_fd
	debug $client_fd "request_method: $request_method"
	# tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*"
	#       / "+" / "-" / "." / "^" / "_" / "`" / "|" / "~"
	#       / DIGIT / ALPHA
	if [[ ! $request_method =~ '^[!#\$%&\x27\*\+\-\.\^_\`|~0-9a-zA-Z]+$' ]]; then
		handle_client_error $client_fd $LINENO "Malformed method"
		return
	fi

	local request_target=
	read -rsd ' ' request_target <&$client_fd
	debug $client_fd "request_target: $request_target"
	# TODO: Parse and classify this
	if [[ ! $request_target =~ "^[!-~]+$" ]]; then
		handle_client_error $client_fd $LINENO "Malformed request-target"
		return
	fi

	local request_http_version=
	local empty=
	read -rsd $CR request_http_version <&$client_fd || true
	read -rsd $LF empty <&$client_fd || true
	if [[ -n "$empty" ]]; then handle_client_error $client_fd $LINENO "Expected LF"; break; fi
	debug $client_fd "request_http_version: $request_http_version"
	if [[ ! $request_http_version =~ "^HTTP/[0-9]\.[0-9]$" ]]; then
		handle_client_error $client_fd $LINENO "Malformed HTTP-version"
		return
	fi
	if [[ ! $request_http_version =~ "^HTTP/1\.[0-9]$" ]]; then
		handle_client_error $client_fd $LINENO "Unsupported HTTP version: $request_http_version"
		return
	fi

	# If this is a TRACE request, handle this special
	if [[ $request_method == "TRACE" ]]; then
		handle_trace $client_fd $request_method $request_target $request_http_version || continue
	fi

	# Read until end of request
	while true; do
		local request_field=
		local empty=
		read -rsd $CR request_field <&$client_fd || true
		read -rsd $LF empty <&$client_fd || true
		if [[ -n "$empty" ]]; then handle_client_error $client_fd $LINENO "Expected LF"; break; fi
		# Stop parsing at first blank line
		if [[ -z "$request_field" ]]; then break; fi
		debug $client_fd "Have field: $request_field ${#request_field}"
	done
	debug $client_fd "Have full request"

	if [[ "$request_target" =~ '^/about:status/' ]]; then
		write_resl $client_fd 200 "OK"
		write_fiel $client_fd "Server" "$SERVER"
		write_fiel $client_fd "Content-Type" "text/plain"
		write_hend $client_fd
		write_line $client_fd "Process request:"
		write_line $client_fd "request_method: $request_method"
		write_line $client_fd "request_target: $request_target"
		write_line $client_fd "request_http_version: $request_http_version"
		write_end $client_fd
		continue;
	fi

	if [[ -n $HANDLER_CGI_TARGET ]]; then
		handle_cgi $client_fd "$HANDLER_CGI_TARGET" "$request_target"
	else
		handle_serve_file $client_fd $request_method $request_target
	fi
}

handle_client_error() {
	local client_fd=$1
	write_resl $client_fd 400 "Client Error"
	write_fiel $client_fd "Server" "$SERVER"
	write_fiel $client_fd "Content-Type" "text/plain"
	write_hend $client_fd
	write_line $client_fd "zhttpd error #$2"
	write_line $client_fd "$3"
	write_end $client_fd
}

# Verify that the request method is one of the allowed methods, return an error if not
# Trace is often used to see how the origin server ends up receiving a request after it's passed through intermediaries
handle_trace() {
	local client_fd=$1
	request_method=$2
	request_target=$3
	request_http_version=$4
	write_resl $client_fd 200 "OK"
	write_fiel $client_fd "Server" "$SERVER"
	write_fiel $client_fd "Content-Type" "application/http"
	write_fiel $client_fd "Allow" "${(j:, :)method_array}"
	write_hend $client_fd
	write_line $client_fd "$request_method $request_target $request_http_version"

	# Read until end of request
	while true; do
		local request_field=
		local empty=
		read -rsd $CR request_field <&$client_fd || true
		read -rsd $LF empty <&$client_fd || true
		write_line $client_fd "$request_field"
		if [[ -n "$empty" ]]; then break; fi
		# Stop parsing at first blank line
		if [[ -z "$request_field" ]]; then break; fi
	done
	write_end $client_fd
	return 1
}

handle_not_found() {
	local client_fd=$1
	write_resl $client_fd 404 "Not Found"
	write_fiel $client_fd "Server" "$SERVER"
	write_fiel $client_fd "Content-Type" "text/plain"
	write_hend $client_fd
	write_line $client_fd "File Not Found: $3"
	write_end $client_fd
}

# Verify that the request method is one of the allowed methods, return an error if not
handle_method() {
	local client_fd=$1
	# request_method is guaranteed to follow the `token` syntax, so it won't contain whitespace or a comma
	local request_method=$2
	local allowed_methods=$3
	local method_array=("${(@s/,/)allowed_methods//[[:space:]]/}")

	# If the request method is in one of the allowed methods, continue without error
	if [[ -n ${(M)method_array:#$request_method} ]]; then
		return 0  # method found
	fi

	# Otherwise print an error
	write_resl $client_fd 405 "Method Not Allowed"
	write_fiel $client_fd "Server" "$SERVER"
	write_fiel $client_fd "Content-Type" "text/plain"
	write_fiel $client_fd "Allow" "${(j:, :)method_array}"
	write_hend $client_fd
	write_line $client_fd "Method not allowed: $request_method"
	write_line $client_fd "Allowed methods: ${(j:, :)method_array}"
	write_line $client_fd ""
	write_end $client_fd
	return 1
}

handle_serve_file() {
	local client_fd=$1
	local request_method=$2
	local request_target=$3

	handle_method $client_fd $request_method "HEAD,GET" || continue

	# Echo back the line with CRLF
	local filepath=$(realpath "$HTROOT$request_target")
	if [[ -z "$filepath" ]]; then
		handle_not_found $client_fd "$filepath"
	elif [[ !  "$filepath" = "$HTROOT"* ]]; then
		handle_client_error $client_fd $LINENO "Not in filepath: $filepath"
	elif [[ -f "$filepath" ]] then
		# If a file, serve the file
		write_resl $client_fd 200 "OK"
		write_fiel $client_fd "Server" "$SERVER"
		# Get the file extension, lowercase
		local ext=${${filepath:e}:l}
		local ct=${ext_mediatype[$ext]:-$DEFAULT_MEDIA_TYPE}
		write_fiel $client_fd "Content-Type" "$ct"
		write_hend $client_fd
		write_file $client_fd "$filepath"
		# Close the listening socket opened by `ztcp -a $listen_fd`
		write_end $client_fd
	elif [[ -d "$filepath" ]] then
		# If a directory, list the directory contents
		write_resl $client_fd 200 "OK"
		write_fiel $client_fd "Server" "$SERVER"
		write_fiel $client_fd "Content-Type" "text/plain"
		write_hend $client_fd
		ls -1 -- "$filepath" >&$client_fd
		write_line $client_fd ""
		write_end $client_fd
	else
		# Ignore other types of files
		# If a directory, list the directory contents
		write_resl $client_fd 403 "Forbidden"
		write_fiel $client_fd "Server" "$SERVER"
		write_fiel $client_fd "Content-Type" "text/plain"
		write_hend $client_fd
		write_line $client_fd "Cannot display file: $filepath"
		write_end $client_fd
	fi
}

# <https://datatracker.ietf.org/doc/html/rfc3875>
# CGI boils down to:
# Read HTTP headers from environment variables
# Write back HTTP/1.0 response with a Status header instead of a status-line
handle_cgi() {
	local client_fd=$1
	local cgi_cmd=$2
	local path_info=$3
	local cgi_response_field=
	local cgi_field_name=
	local cgi_field_value=
	local empty=
	local response_status_code=200
	local response_status_reason="OK"
	declare -A response_headers
	local content_type="text/plain"

	typeset -A cgi_env_vars
	cgi_env_vars=(
		"AUTH_TYPE" ""
		"CONTENT_LENGTH" ""
		"CONTENT_TYPE" ""
		"GATEWAY_INTERFACE" "CGI/1.1"
		"PATH_INFO" ""
		"PATH_TRANSLATED" ""
		"QUERY_STRING" ""
		"REMOTE_ADDR" ""
		"REMOTE_HOST" ""
		"REMOTE_IDENT" ""
		"REMOTE_USER" ""
		"REQUEST_METHOD" ""
		"SCRIPT_NAME" ""
		"SERVER_NAME" ""
		"SERVER_PORT" ""
		"SERVER_PROTOCOL" "HTTP/1.0"
		"SERVER_SOFTWARE" "$SERVER"
	)
	#typeset -a env_pairs
	#for ((i=1; i<=${#cgi_env_vars}; i+=2)); do
	#env_pairs+=("${cgi_env_vars[$i]}=${cgi_env_vars[$i+1]}")
	#done
	#echo $env_pairs env

	while true; do
		while true; do
			read -rsd $CR cgi_response_field || true
			read -rsd $LF empty || true
			if [[ -n "$empty" ]]; then
				handle_client_error $client_fd $LINENO "Expected LF"
				break
			fi
			# Stop parsing at first blank line
			if [[ -z "$cgi_response_field" ]]; then
				break
			fi
			cgi_field_name=${cgi_response_field%%:*}
			cgi_field_value=${cgi_response_field#*:}
			cgi_field_value=${cgi_field_value}
			# Convert field name to lowercase, being case-insensitive
			cgi_field_name=${(L)cgi_field_name}
			# Handle required fields
			if [[ "$cgi_field_name" == "content-type" ]]; then
				content_type="$cgi_field_value"
			elif [[ "$cgi_field_name" == "content-location" ]]; then
				response_headers["Content-Location"]="$cgi_field_value"
			elif [[ "$cgi_field_name" == "status" ]]; then
				# Status line format: "NNN Reason Phrase"
				response_status_code=${cgi_field_value%% *}
				response_status_reason=${cgi_field_value#* }
				# Validate status code is numeric and reasonable
				if [[ ! "$response_status_code" =~ ^[0-9]{3}$ ]] || (( response_status_code < 100 || response_status_code > 599 )); then
					response_status_code=500
					response_status_reason="Internal Server Error"
				fi
			else
				# Store other headers with proper capitalization
				debug $client_fd "$cgi_response_field"
				# Store HTTP header, normalize the field name by lowercasing it
				response_headers[${cgi_field_name:l}]="$cgi_field_value"
			fi
			debug $client_fd "Have CGI response field: $cgi_response_field"
		done
		# Great, we've got the response fields, now send them
		debug $client_fd "Have CGI response body"
		write_resl $client_fd $response_status_code "$response_status_reason"
		write_fiel $client_fd "Content-Type" "$content_type"
		if [[ -n "${response_headers[Content-Location]}" ]]; then
			write_fiel $client_fd "Content-Location" "${response_headers[Content-Location]}"
		fi
		# Write any additional headers
		for header_name in ${(k)response_headers}; do
			[[ "$header_name" != "Content-Location" ]] && write_fiel $client_fd "$header_name" "${response_headers[$header_name]}"
		done
		write_fiel $client_fd "Server" "$SERVER"
		write_hend $client_fd
		# Pass through the body
		cat >&$client_fd
		write_end $client_fd
		break
	done < <(env $cgi_cmd)
	# PHP doesn't like any of these defined and I don't know why
	#done < <(env GATEWAY_INTERFACE=CGI/1.1 PATH_INFO="$path_info" SERVER_SOFTWARE="$SERVER" $cgi_cmd)
	debug $client_fd "Sent off CGI response"
}

# Write response-line
write_resl() {
	local client_fd=$1
	echo -n "HTTP/1.0 $2 $3$CRLF" >&$client_fd
}

# Write header field
write_fiel() {
	local client_fd=$1
	local name=$2
	local value=$3
	# Todo: Verify that $client_fd is in the headers state
	# Todo: Verify that $name follows `token`
	# Todo: Verify that $value follows the header syntax
	echo -n "$name: $value$CRLF" >&$client_fd
}

write_hend() {
	local client_fd=$1
	echo -n "$CRLF" >&$client_fd
}

write_line() {
	local client_fd=$1
	echo -n "$2$CRLF" >&$client_fd
}

write_file() {
	local client_fd=$1
	local file_path=$2
	dd if="$file_path" >&$client_fd
}

write_end() {
	local client_fd=$1
	debug $client_fd "[End]"
	ztcp -c $client_fd
}

debug() {
	local client_fd=$1
	echo "[$1] $2"
}

# Function to display usage
usage() {
	echo "Usage: $0 [options]"
	echo "Options:"
	echo "  -?, --help         Display this help message"
	echo "  --version          Print version"
	echo "  -p, --port=<port>  Listen for HTTP requests on <port>"
	echo "  --cgi=<path>       Forward requests to an invocation of <path>"
	exit 1
}

zparseopts -F '?=opt_help' '-help=opt_help' '-version=opt_version' 'p:=opt_port' '-port:=opt_port' '-cgi:=opt_cgi'  || {
	echo ""
	usage
	exit 1
}

# Handle help option
if [[ -n $opt_help ]]; then
	usage
	exit
fi

# Print server version if --version is specified.
# This is the same string sent in the Server header.
if [[ -n $opt_version ]]; then
	echo $SERVER
	exit
fi

# Trim the leading =, zsh seems to allow --port8080 like a madman
if [[ -n "${opt_port[2]}" && "${opt_port[2]//=}" =~ '^[0-9]+$' ]]; then
	PORT=$MATCH
fi

# Trim --cgi=
if [[ -n "${opt_cgi[2]}" && "${opt_cgi[2]//=}" =~ '^.+$' ]]; then
	HANDLER_CGI_TARGET=$MATCH
fi

# Run the main event loop
# This is just for code aesthetics so that the main loop can be defined
# at the top of the file, and callee functions below it.
main