# Folder structure:
# GPSTools/
#   ├─ pythonTools
#   │   └─ DSD
#   │       ├─ dat64L.py
#   │       ├─ hp8648C.py
#   │       ├─ simple_test_cube_control.py
#   │       └─ simple_chirp_jammer_control.py
#   └─ AntiJamTest
#       ├─ server_main
#       ├─ server_control_spirent
#       ├─ server_control_jammer
#       │   ├─ server.py
#       │   └─ web
#       │       └─ app
#       │           └─ app.py
#       └─ server_control_hackrf

import multiprocessing
import multiprocessing.connection as mp_con
import os, sys
from datetime import datetime
import time

sys.path.append("..\\..\\pythonTools\\DSD")     # For server running on windows
import dat64L
import hp8648C
import simple_test_cube_control
import simple_chirp_jammer_control
import json
import signal

base_dir = './'
sys.path.append(base_dir)

verbose = False

def getChirpStatus(access_lock):
  '''
    Input:
      access_lock = multiprocessing.Manager.Lock() to allow single access to COM port or HackRF, otherwise there will be permission issue
    Return:
      dictionary of serial status
  '''
  fn = 'getChirpStatus()'
  d = {}

  num_try = 2
  acq_res = None
  while num_try > 0:
    if verbose: print(fn,"Acquiring access_lock...",num_try)
    acq_res = access_lock.acquire(False) 
    if acq_res:
      break
    time.sleep(0.5)
    num_try -= 1
  
  if acq_res:
    if verbose: print(fn,"Acquired access_lock.")
    d['status'] = 'Not in use'
    access_lock.release()
    if verbose: print(fn,"Released access_lock.")
  else:
    if verbose: print(fn,"Cannot acquire access_lock.")
    try:
      with open("chirp_status.json") as f: 
        d = json.load( f )
    except:
      d['status'] = 'In use'
  return d


def getdat64LStatus(dat64_port, access_lock):
  '''
    Input:
      dat64_port = int, eg. 3,4, or 19, for COM3, COM4, or COM19
      access_lock = multiprocessing.Manager.Lock() to allow single access to COM port, otherwise there will be permission issue when access COM at the same time.
    Return:
      dictionary of attenuator status
  '''
  fn = 'getdat64LStatus(dat64_port)'
  dat64_port = int(dat64_port)
  d = {}

  num_try = 2
  acq_res = None
  while num_try > 0:
    if verbose: print(fn,"Acquiring access_lock...",num_try)
    acq_res = access_lock.acquire(False) 
    if acq_res:
      break
    time.sleep(0.5)
    num_try -= 1

  if acq_res:
    if verbose: print(fn,"Acquired access_lock.")
    attenuation_dB = dat64L.get_attenuation_dB(dat64_port)
    d['attenuation_dB'] = attenuation_dB.replace('\r','').rstrip()
    access_lock.release()
    if verbose: print(fn,"Released access_lock.")
  else:
    if verbose: print(fn,"Cannot acquire access_lock.")
    d['attenuation_dB'] = 'Try again.'
  
  return d

def getHP8648CStatus(gdib_address):
  '''
  Input: gdib_address = int, GPIB port address eg. 16 or 18
  Return: dictionary of HP8648C status
  '''
  fn = 'getHP8648CStatus(gdib_address)'
  GPIB_addr = int(gdib_address)
  d = {}

  try:
    sigGen = hp8648C.device_connect(GPIB_addr)
    freq = sigGen.get_freq_MHz()
    pwr = sigGen.get_power_dBm()
    on = sigGen.is_rf_on()
    sigGen.disconnect()

    d['on'] = on
    d['freq_MHz'] = freq
    d['power_dBm'] = pwr
  except:
    d['msg'] = 'Failed'

  return d


def set_cw(GPIB=18, pwr=-999.0, freq=1580.42, access_lock=None):
  '''
  Input:  GPIB = int, gpib port
          pwr = float, power level, -999.0 means turning off
          freq = float, at which frequency to increase power
          access_lock = multiprocessing.Manager.Lock() to allow single access to COM for attenuator
  Return: True/False
  '''
  fn = 'set_cw()'
  if verbose: print(fn,':: set GPIB,desired power,frequency:', GPIB, pwr, freq)
  
  num_try = 3
  acq_res = None
  while num_try > 0:
    if verbose: print(fn,"Acquiring access_lock...",num_try)
    acq_res = access_lock.acquire(False) 
    if acq_res:
      break
    time.sleep(0.5)
    num_try -= 1
  
  if acq_res:
    if verbose: print(fn,"Acquired access_lock. call set()")
    simple_test_cube_control.set(GPIB, pwr, freq)
    access_lock.release()
    if verbose: print(fn,"Released access_lock.")
  else:
    if verbose: print(fn,"Cannot acquire access_lock.")
    return False

  return True

def wrap_turn_on_chirp(access_lock, com_port, desired_pwr_dBm, sec):
  '''
  This function needs to be defined at the top-level to avoid Attribute Error: can't pickle local object
  Functions are only picklable if they are defined at the top-level of a module.
  Nested functions defind within turn_chirp_jammer_on() won't be importable by the child.
  '''
  fn = 'wrap_turn_on_chirp()'
  simple_chirp_jammer_control.set(com_port, desired_pwr_dBm, sec)
  access_lock.release()
  if verbose: print(fn,"Released access_lock.")
  
def turn_chirp_jammer_on(com_port, desired_pwr_dBm, sec, access_lock):
  '''
  This function returns after simple_chirp_jammer_control.set() finishes
  Input:  com_port = int
          desired_pwr_dBm = float, power level
          sec = float, peirod of chirp is set
          access_lock = multiprocessing.Manager.Lock() to allow single access to chirp
  Return: True/False
  '''

  fn = 'turn_chirp_jammer_on()'

  num_try = 1
  acq_res = None
  response = {}

  while num_try > 0:
    if verbose: print(fn,"Acquiring access_lock...", num_try)
    acq_res = access_lock.acquire(False) 
    if acq_res:
      break
    time.sleep(0.1)
    num_try -= 1
  
  if acq_res:
    if verbose: print(fn,"Acquired access_lock. call set()")
    proc = multiprocessing.Process( target=wrap_turn_on_chirp, args=(access_lock, com_port, desired_pwr_dBm, sec,) )
    proc.start()
    end_time = datetime.fromtimestamp(time.time()+sec).strftime('%Y-%m-%d %H:%M:%S')
    response = {'status': 'In use', 'end_time':end_time, 'pid':proc.pid}
    with open("chirp_status.json","w") as f:
      json.dump( response, f )
    return response
  else:
    if verbose: print(fn,"Cannot acquire access_lock.")
    try:
      with open("chirp_status.json") as f: 
        response = json.load( f )
    except:
      response = {'status': 'In use'}
    return response


def turn_chirp_jammer_off(access_lock):
  '''
  If lock is locked, force to kill process by reading pid from chirp_status.json
  Return: True/False
  '''

  fn = 'turn_chirp_jammer_off()'

  num_try = 1
  acq_res = None
  response = {}

  while num_try > 0:
    if verbose: print(fn,"Acquiring access_lock...", num_try)
    acq_res = access_lock.acquire(False) 
    if acq_res:
      break
    time.sleep(0.1)
    num_try -= 1
  
  if acq_res:
    # Do nothing
    access_lock.release()
    if verbose: print(fn,"Released access_lock.")
    return False
  else:
    if verbose: print(fn,"Cannot acquire access_lock.")
    try:
      with open("chirp_status.json") as f: 
        d = json.load( f )
      pid = int(d['pid'])
      if verbose: print(fn,"pid",pid)
      os.kill(pid, signal.SIGTERM)
      if verbose: print(fn,"Kill all process.")
    except:
      if verbose: print(fn,"Load chirp status failed.")

    if verbose: print(fn,"Release lock anyway.")
    access_lock.release()
    return True

def handle_req(conn, lck_cw, lck_chirp):
  '''
  Input: conn = socket for communication with app.py
         lck_cw = multiprocessing.Manager.Lock() allows single access to COM for attenuator used by CWs. Racing occurs when accessing COM ports.
         lck_chirp = multiprocessing.Manager.Lock() allows single access to COM for chirp + attenuator.
  Handle requests from app.py and send a response back over the 'conn' socket.
  '''
  msg = conn.recv()
  if verbose: 
    print("Got msg:", msg)
  if msg[0] == 'getHP8648CStatus':
    if verbose: print(msg)
    rt = getHP8648CStatus(int(msg[1]))
    conn.send( rt )
  elif msg[0] == 'getDat64LStatus':
    if verbose: print(msg)
    if int(msg[1]) == 19:
      rt = getdat64LStatus(int(msg[1]), lck_chirp) # attenuator@COM19 is used by Chirp@COM20
      if rt['attenuation_dB'] == 'Try again.':
        rt['attenuation_dB'] = 'Unknown'
        rt['status'] = 'In use'
    else:
      rt = getdat64LStatus(int(msg[1]), lck_cw) # attenuator@COM3,COM4 is used by CWs
    conn.send( rt )
  elif msg[0] == 'getChirpStatus':
    if verbose: print(msg)
    rt = getChirpStatus(lck_chirp)
    conn.send( rt )
  elif msg[0] == 'turnOnCW':
    if verbose: print(msg)
    rt = set_cw( int(msg[1]), float(msg[2]), float(msg[3]), lck_cw)
    if verbose: print("send(rt=",rt,")")
    conn.send( rt )
  elif msg[0] == 'turnOnChirp':
    if verbose: print(msg)
    rt = turn_chirp_jammer_on( int(msg[1]), float(msg[2]), float(msg[3]), lck_chirp)
    conn.send( rt )
  elif msg[0] == 'turnOffChirp':
    if verbose: print(msg)
    rt = turn_chirp_jammer_off( lck_chirp )
    conn.send( str(rt) )
  else:
    print("Unknown msg:",msg)
    d = {}
    d['msg'] = 'Unknown msg:' + msg
    conn.send( d )
  conn.close()
  log_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f') + ": " + " ".join(map(str,msg))
  if verbose: print(log_str)

def run_server(lck_cw, lck_chirp):
  '''
  Input: 
    lck_cw = multiprocessing.Manager.Lock() to allow single access to COM for attenuator COM 3,4
    lck_chirp = multiprocessing.Manager.Lock() to allow single access to COM 19, 20
  Start listening on port ? for requests from app.py.  Each request is
  handled in a separate thread so multiple people can access the server.
  '''
  # Using port 10001 might lead to failure for Windows Task Scheduler to auto launch at startup.
  control_port = 10001

  addr = ('localhost', control_port)
  print('Listen on {}'.format(addr))
  
  now = datetime.now()
  dt_string = now.strftime("%d/%m/%Y %H:%M:%S")

  listener = mp_con.Listener(addr)
  try:
    while True:
      conn = listener.accept()
      proc = multiprocessing.Process(target=handle_req, args=(conn, lck_cw, lck_chirp))
      proc.start()
      # conn.close()
  except KeyboardInterrupt:
    print("Receive Ctrl + C")
    pass

if __name__ == "__main__":

  now = datetime.now()
  dt_string = now.strftime("%d/%m/%Y %H:%M:%S")
  f = open("start.log", 'a')
  print("Date and time when launching the server =", dt_string)
  print(os.getcwd())
  f.write(dt_string+'\n')
  f.write(os.getcwd()+'\n')
  f.close()

  os.environ['no_proxy'] = '*'
  m = multiprocessing.Manager()
  lck_cw = m.Lock()
  lck_chirp = m.Lock()
  run_server(lck_cw, lck_chirp)