open Unix;;

exception No_home_variable;;
exception No_home_to_chdir;;
exception Fails_forking;;
exception Already_running;;
exception No_pidfile;;
exception Cant_find_logger;;
exception Problem_with_logger;;

type daemon_state = Terminate | Respawn | MlnetDied | Nop

type action_type = Start | Stop | Restart 

type mldonkey_server_state = {
	mutable daemon : bool;
	mutable quiet : bool;
	mutable logfile : string option;
	mutable pidfile : string option;
	mutable program : string;
	mutable logger : string;
	mutable prg_args : string list;
	mutable chdir : string option;
	mutable chuid : int * int;
	mutable umask : int;
	mutable action : action_type;

	(* Variable to determine the way we manage mlnet *)
	mutable end_duration : float;
	mutable max_alive : float;
	mutable min_alive : float;
	mutable respawn_after : float;
}

let user_home () =
	try 
		getenv "HOME"
	with Not_found ->
		raise No_home_variable

let get_chdir st =
	match st.chdir with
	None ->
		Filename.concat (user_home ()) ".mldonkey"
	| Some x ->
		x

let get_pidfile st =
	match st.pidfile with
	None ->
		Filename.concat (get_chdir st) "mldonkey.pid"
	| Some x ->
		x

let get_logfile st =
	match st.logfile with
	None ->
		Filename.concat (get_chdir st) "mldonkey.log"
	| Some x ->
		x
	
let create_mldonkey_server_state = 
	{
	daemon = false;
	quiet = false;
	logfile = None;
	pidfile = None;
	program = "/usr/bin/mlnet";
	logger = "/usr/bin/logger";
	prg_args = [];
	chdir = None;
	chuid = (Unix.getuid (), Unix.getgid ());
	umask = 0o0022;
	action = Start;

	(* 5 min *)
	end_duration = 300.;

	(* 15 min *)
	min_alive = 900.;

	(* 1 day *)
	max_alive = 88400.;

	(* 10 min *)
	respawn_after = 600.;
	}


let get_mldonkey_server_state () =
	let state = create_mldonkey_server_state
	in
	let add_args x = 
		state.prg_args <- x :: state.prg_args
	in
	let _ = Arg.parse [
	("--daemon", Arg.Unit ( fun x -> state.daemon <- true ), 
		"Run in daemon mode");
	("--quiet", Arg.Unit ( fun x -> state.quiet <- true ), 
		"Keep quiet");
	("--logfile", Arg.String ( fun x -> state.logfile <- Some x ), 
		"Where to put the log");
	("--pidfile", Arg.String ( fun x -> state.pidfile <- Some x ), 
		"Where to put the pid");
	("--program", Arg.String ( fun x -> state.program <- x ), 
		"Which program to start");
	("--logger", Arg.String ( fun x -> state.logger <- x ),
		"Logger program");
	("--chdir", Arg.String ( fun x -> state.chdir <- Some x ), 
		"Where to chdir");
	("--chuid", Arg.String ( fun x -> 
		let passwd_ent =
			Unix.getpwnam x
		in
			state.chuid <- (passwd_ent.pw_uid, passwd_ent.pw_gid)
		), 
		"Who own the process");
	("--umask", Arg.String ( fun x ->
		state.umask <- int_of_string ( "0o"^x ) 
		),
		"What umask to use");
	("--end-duration", Arg.Int ( fun x -> state.end_duration <- float_of_int x ),
		"How much time does it take to end mlnet");
	("--max-alive", Arg.Int ( fun x -> state.max_alive <- (float_of_int x) *. 3600. ),
		"For how long an instance of mlnet should run");
	("--min-alive", Arg.Int ( fun x -> state.min_alive <- float_of_int x),
		"Minimun time between respawning");
	("--respawn-after", Arg.Int ( fun x -> state.respawn_after <- float_of_int x ),
		"When mlnet fails, how long to wait before restarting");
	("--start", Arg.Unit ( fun x -> state.action <- Start),
		"Start mldonkey_server");
	("--stop", Arg.Unit ( fun x -> state.action <- Stop),
		"Stop a running mldonkey_server ( use the pidfile )");
	("--restart", Arg.Unit ( fun x -> state.action <- Restart),
		"Restart a running mldonkey_server ( use the pifile, only respawn mlnet )");
	("--", Arg.Rest ( fun x -> add_args x ),
		"MLnet arguments")
	]
	add_args
	"Usage mldonkey_server [options] -- [mlnet options] where options are :"
	in
	state

let debug st str =
	if st.quiet then
		()
	else if st.daemon then
		match Unix.system (st.logger^" -t mldonkey_server \""^str^"\"") with
		WEXITED(0) ->
			()
		| WEXITED(127) ->
			raise Cant_find_logger
		| _ ->
			raise Problem_with_logger
	else
	begin
		print_string str;
		print_newline ()
	end

let go_home st =
	debug st ("Chdir to home dir : "^(get_chdir st));
	Sys.chdir (get_chdir st)

let create_home st =
	if Sys.file_exists (get_chdir st) then
		()
	else
		if not st.daemon then
		begin
			let answer =		
				prerr_string ((get_chdir st)^" doesn't exists."^
				" Do you want to create it ? ( y/N )");
				flush(Pervasives.stderr);
				read_line ()
			in
			match answer with
			"y" ->
				Unix.mkdir (get_chdir st) 0o0755;
				debug st ("Creating home dir : "^(get_chdir st))
			| _ ->
				raise No_home_to_chdir
		end
		else
			raise No_home_to_chdir

let check_no_other_mldonkey_server st =
	if Sys.file_exists (get_pidfile st) then
		raise Already_running
	else
		()

let set_uid st =
	let (uid,gid) = st.chuid
	in
	debug st ("Set uid/gid of the process ( "^(string_of_int uid)^", "^(string_of_int gid)^" )");
	setgid gid;
	setuid uid

let set_umask st =
	debug st ("Set umask of the process :"^(string_of_int st.umask));
	ignore(umask st.umask)

let create_pidfile st =
	let pidfile = open_out (get_pidfile st)
	in
	debug st ("Writing PID ( "^(string_of_int (Unix.getpid ()))^" ) to pidfile : "^(get_pidfile st));
	output_string pidfile (string_of_int (Unix.getpid ()));
	output_string pidfile "\n";
	close_out pidfile

let read_pidfile st =
	try
		let pidfile = open_in (get_pidfile st)
		in
		let pid_server = int_of_string (input_line pidfile)
		in
		debug st ("Reading PID ( "^(string_of_int pid_server)^" from pidfile : "^(get_pidfile st));
		close_in pidfile;
		pid_server
	with Sys_error(_) ->
		raise No_pidfile

let close_pidfile st =
	try
		Sys.remove (get_pidfile st)
	with Sys_error(_) ->
		raise No_pidfile

let stop_or_die st pid =
	let timeout = ref false
	in
	let _ =
		Sys.set_signal Sys.sigalrm (Sys.Signal_handle
			( fun x -> timeout := true ));
		ignore (Unix.alarm (int_of_float st.end_duration));
		debug st ("Waiting termination of process "^(string_of_int pid));
		try 
			Unix.kill pid Sys.sigterm;
			ignore (waitpid [] pid )
		with Unix.Unix_error(_, _, _) ->
			()
	in
	if !timeout then
	begin
		debug st ("Process "^(string_of_int pid)^" not responding, taking measure : SIGKILL");
		try
			Unix.kill pid Sys.sigkill
		with Unix.Unix_error(_, _, _) ->
			()
	end
	else
		debug st ("Process "^(string_of_int pid)^" terminated")

let daemonize st = 
	if st.daemon then
		if Unix.fork () = 0 then
			if Unix.setsid () = Unix.getpid () then
				if Unix.fork () = 0 then
					()
				else
					exit 0
			else
				raise Fails_forking 
		else
			exit 0
	else
		()	


let start_mldonkey_server st =
	set_uid st;
	set_umask st;
	create_home st;
	go_home st;
	check_no_other_mldonkey_server st;
	daemonize st;
	create_pidfile st;
	let launch_mlnet st =
		let (logger_stderr, mlnet_stderr) =
			if st.daemon then
				Unix.pipe ()
			else
				(stdin, stderr)
		in
		let (logger_stdout, mlnet_stdout) =
			if st.daemon then
				Unix.pipe ()
			else
				(stdin, stdout)
		in
		let args = Array.of_list 
			(st.program :: (List.rev st.prg_args))
		in
		let pid_mlnet = 
			debug st ("Launching MLnet process");
			create_process st.program args
			stdin
			mlnet_stdout
			mlnet_stderr
		in
		let pid_logger_stderr =
			if logger_stderr != stdin then
			begin
				debug st ("Launching MLnet stderr logger");
				create_process st.logger [| st.logger ; "-t"; "mlnet_error" |]
				logger_stderr
				stdout
				stderr
			end
			else
				0
		in
		let pid_logger_stdout =
			if logger_stdout != stdin then
			begin
				debug st ("Launching MLnet stdout logger");
				create_process st.logger [| st.logger ; "-t"; "mlnet" |]
				logger_stdout
				stdout
				stderr
			end
			else
				0
		in
		(
			[pid_mlnet; pid_logger_stderr; pid_logger_stdout],
			[logger_stderr; mlnet_stderr; logger_stdout; mlnet_stdout]
		)
	in
	let stop_mlnet st (pids, fds) =
		let str_pids = List.fold_left 
			( fun str x -> str^" "^(string_of_int x) ) 
			"" pids
		in
		debug st ("Stopping processes PID ("^str_pids^" )") ;
		begin
		try
			let close_fds x =
				if x != stdout && x != stdin && x != stderr then
					Unix.close x
				else
					()
			in
			let stop_pids x =
				if x != 0 then
					stop_or_die st x
				else
					()
			in
			List.iter stop_pids pids;
			List.iter close_fds fds
		with Unix.Unix_error(_,_,_) ->
			()
		end;
		debug st ("Process stopped PID ("^(str_pids)^" )")
	in
	let state = ref Nop
	in
	let terminate = ref false
	in
	let reload = ref false
	in
	let _ =
		Sys.set_signal Sys.sigint ( Sys.Signal_handle 
			( fun x -> state := Terminate ));
		Sys.set_signal Sys.sigterm ( Sys.Signal_handle
			( fun x -> state := Terminate ));
		Sys.set_signal Sys.sighup ( Sys.Signal_handle
			( fun x -> state := Respawn ));
		Sys.set_signal Sys.sigchld ( Sys.Signal_handle
			( fun x -> state := MlnetDied ));
		Sys.set_signal Sys.sigpipe ( Sys.Signal_handle
			( fun x -> state := Terminate ))
	in
	while not !terminate do
		let mlnet = launch_mlnet st 
		in
		let last_respawn = Unix.time ()
		in
		Sys.set_signal Sys.sigalrm ( Sys.Signal_handle
			( fun x -> state := Respawn ));
		ignore (Unix.alarm ( int_of_float st.max_alive ));
		reload := false;
		while not !reload && not !terminate do
			let _ = 
				try
					state := Nop;
					begin
					match Unix.wait () with
					x,WEXITED(y) ->
						debug st ("Process PID ( "^(string_of_int x)^
						" ) exit with return code "^(string_of_int y))
					| x,WSIGNALED(y) ->
						debug st ("Process PID ( "^(string_of_int x)^
						" ) was killed by signal "^(string_of_int y))
					| x,WSTOPPED(y) ->
						debug st ("Process PID ( "^(string_of_int x)^
						" ) was stopped by signal "^(string_of_int y))
					end;
					state := MlnetDied
				(* On peut etre interrompu par un signal extérieur *)
				with Unix.Unix_error(EINTR,_,_) ->
					()
			in
			match !state with
			Terminate ->
				debug st ("Terminate process");
				stop_mlnet st mlnet;
				terminate := true
			| Respawn ->
				debug st ("Respawn process");
				stop_mlnet st mlnet;
				reload := true
			| MlnetDied ->
				if Unix.time () -. last_respawn < st.min_alive then
				begin
					debug st ("Process respawning too fast : only live "^
						(string_of_float(Unix.time () -. last_respawn)));
					stop_mlnet st mlnet;
					terminate := true
				end
				else
				begin
					debug st ("Process died, respawning : live for "^
						(string_of_float(Unix.time () -. last_respawn)));
					stop_mlnet st mlnet;
					reload := true
				end
			| Nop ->
				()
		done;
	done;
	debug st "MLDonkey server end";
	close_pidfile st


let kill_mldonkey_server st signal signal_name=
	let pid_server = read_pidfile st
	in
	debug st ("Sending signal "^signal_name^" to process PID ( "^(string_of_int pid_server)^" )");
	if signal = Sys.sigterm then
		stop_or_die st pid_server
	else
	begin
		try
			Unix.kill pid_server signal
		with Unix.Unix_error (_, _, _) ->
			()
	end

let stop_mldonkey_server st =
	kill_mldonkey_server st Sys.sigterm "SIGTERM"

let restart_mldonkey_server st =
	kill_mldonkey_server st Sys.sighup "SIGHUP"

let _ = 
	try
		let state = get_mldonkey_server_state ()
		in
		begin
		match state.action with
		Start ->
			start_mldonkey_server state
		| Stop ->
			stop_mldonkey_server state
		| Restart ->
			restart_mldonkey_server state
		end;
		exit 0
	with No_home_variable ->	
		begin
		prerr_string ("Could not guess $HOME environnement variable : provide a --chdir or $HOME");
		prerr_newline ()
		end
	| No_home_to_chdir ->
		begin
		prerr_string ("Home dir doesn't exist");
		prerr_newline ()
		end
	| Fails_forking ->
		begin
		prerr_string ("Cannot fork process");
		prerr_newline ()
		end
	| Already_running ->
		begin
		prerr_string ("Some others mldonkey_server are running ( a pidfile exists )");
		prerr_newline ()
		end
	| No_pidfile ->
		begin
		prerr_string ("No pidfile, maybe no mldonkey_server are running");
		prerr_newline ()
		end
	| Unix.Unix_error (error,_,_) ->
		begin 
		prerr_string (error_message error);
		prerr_newline ()
		end