#!/usr/bin/env python

import os
import sys
import signal
import time
from socket import gethostname
from subprocess import Popen,PIPE
debug=0

class daemon:
	def __init__(self):
		self.pid_filename='/var/run/nvtempd.pid'
		self.modes={0:'Auto',1:'Manual'}
		self.default_mode=0

		################### Editable bits #################
		self.max_fan_speed=80 #hard coded maximum by nvidia YMMV (in%)
		self.min_fan_speed=30 #hard coded minimum by nvidia YMMV (in%)
		self.go_auto_temp=40  #below this temperature we let the card take care of itself
		self.fans=[0,1,2]         #self.fans=[0,1,2,3] #add more if you have more
		self.exe='/usr/bin/nvidia-settings'
		self.ENV={'DISPLAY':':0'}
		self.gpu=0
		################### End Editable bits #############

		self.match="(%s%s[gpu:%s])"%(gethostname(),self.ENV['DISPLAY'],self.gpu)
		self.firstrun=1
		self.curve={}
		x=0
		while x <= 120: #GPUCoreTemp only reposts to the nearest degree C, no point running the curve every time, cheaper to make a dictionary of temp vs fan speed.
			#self.curve.update({x:(((x- a	)/  b )** c + d)  alter a &| b to move the curve left/right , c to set the steepness of the change, can be visualized in kmplot with y=(((x-15)/25.7)^5)+30 
			self.curve.update({x:(((x-15)/25.7)**5)+30}) #
			if self.curve[x]>self.max_fan_speed:self.curve.update({x:self.max_fan_speed})
			if self.curve[x]<self.min_fan_speed:self.curve.update({x:self.min_fan_speed})
			x+=1

	def exit(self,sig,*_):
		sig_name=str(signal.Signals(sig))[8:]
		print('NVtempd: Received %s, stopping gracefully ...' %sig_name)
		self.set_mode(self.default_mode)
		try:os.remove(self.pid_filename)
		except:sys.exit(1)
		sys.exit(0)

	def get_gpu_temp(self,gpu=0):
		child=Popen([self.exe, '-q', '[gpu:%s]/GPUCoreTemp'%self.gpu], env=self.ENV, stdout=PIPE)
		child.wait()
		for data in child.stdout.readlines():
			line=str(data)
			if line.find(self.match) != -1:
				break
		return int(line.strip()[line.rfind(' ')+1:-4])

	def start(self):
		self.get_mode()
		try:
			pid_file = open(self.pid_filename,'x')
		except FileExistsError as error:
			pid=self.read_file(self.pid_filename)
			if not os.system('ps -A|awk \'{print $1}\'|grep -qw "%s"'%pid):
				print('NVtempd: PID file already exists and daemon is running. Exiting...')
				sys.exit(1)
			else:
				print('NVtempd: Removing stale PID file:')
				os.remove(self.pid_filename)
				self.start()
		except OSError as error:
			print(error)
			sys.exit(1)
		#endtry
		pid_file.write(str(os.getpid()))
		pid_file.close()
		signal.signal(signal.SIGINT, self.exit)#from ctrl-c
		signal.signal(signal.SIGUSR1, self.exit)#from main script stop()
		signal.signal(signal.SIGTERM, self.exit)#from start-stop-daemon
		while True:
			time.sleep(1)
			self.check_temp()

	def check_temp(self):
		if self.firstrun:
			self.set_mode(1)
			if debug: print('making some noise')
			self.set_fan_speed(self.max_fan_speed)
			time.sleep(5)
			self.set_fan_speed(self.min_fan_speed)
			self.set_mode(0)
			self.firstrun=0
		if debug: print("checking temp")
		gpu_temp=self.get_gpu_temp(self.gpu)
		if debug: print("temp=%s"%gpu_temp)
		if gpu_temp <= self.go_auto_temp:
			if self.current_mode != 'Auto':
				if debug:print("going auto")
				self.set_mode(0)
			return
		self.set_fan_speed(self.curve[gpu_temp])

	def set_mode(self,mode):
		child=Popen([self.exe, '-a', '[gpu:%s]/GPUFanControlState=%s'%(self.gpu,mode)],env=self.ENV, stdout=PIPE, stderr=PIPE)
		child.wait()
		if debug:
			print(child.stdout.readlines())
			print(child.stderr.readlines())
		self.current_mode=self.modes[int(mode)]

	def set_fan_speed(self,Pspeed):
		self.get_mode()
		if self.current_mode != 'Manual':
			if debug:print("going manual")
			self.set_mode(1)
		speed=str(int(Pspeed))
		for fan in self.fans:
			child=Popen([self.exe, '-a', '[fan:%s]/GPUTargetFanSpeed=%s'%(fan,speed)],env=self.ENV, stdout=PIPE, stderr=PIPE)
			child.wait()

	def get_fan_speed(self,fan):
		child=Popen([self.exe, '-q', '[fan:%s]/GPUTargetFanSpeed'%(fan)],env=self.ENV, stdout=PIPE, stderr=PIPE)
		child.wait()
		match="(%s%s[fan:%s])"%(gethostname(),self.ENV['DISPLAY'],fan)
		rpm=0
		for data in child.stdout.readlines():
			line=str(data)
			if line.find(match) != -1:
				rpm=int(line.strip()[line.rfind(' ')+1:-4])
				break
		return rpm

	def get_rpm(self,fan):
		child=Popen([self.exe, '-q', '[fan:%s]/GPUCurrentFanSpeedRPM'%fan], env=self.ENV, stdout=PIPE, stderr=PIPE)
		child.wait()
		match="(%s%s[fan:%s])"%(gethostname(),self.ENV['DISPLAY'],fan)
		rpm=0
		for data in child.stdout.readlines():
			line=str(data)
			if line.find(match) != -1:
				rpm=int(line.strip()[line.rfind(' ')+1:-4])
				break
		return rpm

	def get_mode(self):
		child=Popen([self.exe, '-q', '[gpu:%s]/GPUFanControlState'%self.gpu], env=self.ENV, stdout=PIPE)
		child.wait()
		for data in child.stdout.readlines():
			line=str(data)
			if line.find(self.match) != -1:
				mode=int(line.strip()[line.rfind(' ')+1:-4])
				break
		self.current_mode=self.modes[mode]

	def read_file(self,path):
		try:
			ret_file=open(path,'r')
			ret_val=ret_file.read()
			ret_file.close()
			return int(ret_val)
		except:
			print("read file error")
			self.exit(signal.SIGUSR1)

	def simple_exit(self,sig,*frame):
		sys.exit(0)

	def force_stop(self,pid):
		if pid:print('NVtempd: Daemon refused to exit cleanly, killing PID:%s with SIGKILL and cleaning up'%pid)
		try: os.remove(self.pid_filename)
		except: pass
		if pid:
			print ('NVtempd: %s'%os.system('/bin/kill -SIGKILL %s'%pid))
		else:
			os.system('ps aux|sed -n "/python %s start/p"|awk \'{print $2}\'|xargs /bin/kill -SIGKILL >/dev/null 2>&1'%(sys.argv[0].replace('/','\\/')))
		sys.exit(0)

	def stop(self):
		pid=0
		try:
			pid=self.read_file(self.pid_filename)
			os.remove(self.pid_filename)
		except:
			self.force_stop(pid)
		os.system('/bin/kill -SIGUSR1 %s'%pid)
		timer=60
		while timer > 0 and not os.system('ps -A|grep -qw "%s"'%pid):
			time.sleep(0.1)
			timer-=1
		if not os.system('ps -A|grep -q "%s"'%pid):self.force_stop(pid)
		sys.exit(0)

	def status(self):
		signal.signal(signal.SIGINT,self.simple_exit)
		try: interval=sys.argv[2]
		except:interval=0
		run=True
		while run:
			if os.system('ps aux|grep -v grep 2>/dev/null|grep -q "python %s start"'%sys.argv[0]):
				isRunning='Not running'
			else:
				isRunning='Yes, PID:%s'%self.read_file(self.pid_filename)
			temp=self.get_gpu_temp(self.gpu)
			self.get_mode()
			print('#####################################')
			print(' Running            : %s'%isRunning)
			print(' Temp               : %s°C'%temp)
			for fan in self.fans:
				print(' Fan %s Set to       : %s%%'%(fan,self.get_fan_speed(fan)))
				print(' Fan %s Current RPM  : %s'%(fan,self.get_rpm(fan)))
			print(' Current Mode       : %s'%self.current_mode)
			print('#####################################')
			time.sleep(int(interval))
			run=interval


def usage():
	print('Usage: %s <start|stop|status [update interval]>'%sys.argv[0][sys.argv[0].rfind('/')+1:])

try:
	action=sys.argv[1]
	if action == 'start':
		main=daemon()
		if debug:
			main.start()
			sys.exit()
		fork_ret=os.fork()
		if fork_ret == 0:
			main.start()
		else:
			sys.exit(0)
	elif action == 'stop':
		main=daemon()
		main.stop()
	elif action == 'status':
		main=daemon()
		main.status()
	else:
		usage()
except IndexError:
	usage()
#endtry