# cimport C libraries
cimport cython
cimport numpy as np
from libc.stdio cimport printf, stdout, fflush, remove
from libc.stdio cimport printf, snprintf, stdout, fflush, remove
from libc.stdlib cimport exit, malloc, free, EXIT_SUCCESS, EXIT_FAILURE
from libc.stdint cimport int8_t, uint8_t, int16_t, uint16_t, int32_t, uint32_t, int64_t, uint64_t
from libc.signal cimport SIGINT, SIGUSR1, SIGALRM, SIGBUS, SIGSEGV, SIGQUIT, SIGUSR2
from posix.signal cimport kill, sigaction, sigaction_t, sigset_t, sigemptyset, sigaddset, sigfillset, kill
from posix.unistd cimport getppid, pause, close, getpid, write
from posix.mman cimport shm_open, shm_unlink, mmap, munmap, PROT_READ, PROT_WRITE, MAP_SHARED, mlockall, MCL_CURRENT, MCL_FUTURE, munlockall
from posix.fcntl cimport O_RDWR, O_WRONLY, O_CREAT, O_TRUNC, open
from posix.types cimport pid_t
from libc.errno cimport errno, EINTR, EPIPE
from libc.string cimport memset, strcpy, strcat, strlen, memcpy
from posix.time cimport clock_gettime, CLOCK_MONOTONIC_RAW, timespec
from libcpp cimport bool
DEF SIG_BLOCK = 1

# import Python libraries
import numpy as np
{% if is_writer -%}
import SharedArray as sa
{%- endif %}
import time

# headers for all sinks
cdef extern from "utilityFunctions.h" nogil:
  enum: __GNU_SOURCE
  void die(char *errorStr)
  void open_shared_mem(uint8_t **ppmem, const char *pName, int numBytes, int shm_flags, int mmap_flags)
  void set_sighandler(int signum, void *psh, sigset_t *block_mask)
  void make_realtime()

cdef extern from "utilityFunctions.h":
  void init_utils(void (*pHandleExit)(int exitStatus), sigset_t *pExitMask)

cdef extern from "semaphore.h" nogil:
  cdef union sem_union:
    pass
  ctypedef sem_union sem_t
  sem_t *sem_open(const char *name, int oflag)
  int sem_wait(sem_t *)
  int sem_post(sem_t *)
  int sem_close(sem_t *sem)

cdef extern from "constants.h":
  enum: NUM_CHANNELS
  enum: MAX_PATH_LEN
  enum: NUM_CHILDREN
  enum: PAGESIZE
  enum: SMEM0_HEADER_LEN
  cdef const char *SMEM0_PATHNAME
  cdef const char *SMEMSIG_PATHNAME
  enum: NUM_NON_SOURCES
  enum: NUM_SEM_SIGS
  enum: NUM_INTERNAL_SIGS
  size_t ROUND_UP(int X, int Y)
  enum: BUF_VARS_LEN
  enum: SHM_SIZE
  enum: NUM_TICKS_OFFSET
  enum: BUF_VARS_OFFSET
  enum: ASYNC_WRITER_MUTEXES_OFFSET
  enum: SEM_NAME_LEN

{%- if async %}
cdef extern from "<pthread.h>" nogil:
  struct sched_param:
    int sched_priority
{%- endif %}

{%- if is_writer %}
{{ driver_code["imports"] }}
{%- endif %}

# cdef variables
cdef pid_t pid, ppid
cdef sigset_t exitMask
cdef char pathName[MAX_PATH_LEN]

cdef size_t shmSize
cdef uint8_t *pmem
cdef int64_t *pNumTicks
{% if is_main_process -%}
cdef sem_t *pTickUpSem
cdef sem_t *pTickDownSem
cdef sem_t *pSigSems[NUM_SEM_SIGS]
cdef char semNameBuf[SEM_NAME_LEN]
{%- endif %}
cdef uint32_t *pBufVars
{%- for sig, dtype in sig_types.items() %}
cdef uint32_t *{{sig}}BufVars
cdef {{dtype}} *{{sig}}StrtPtr
cdef {{dtype}} *{{sig}}EndPtr
{%- endfor %}

{% if is_writer -%}
cdef uint32_t packetSize = {{out_signal_size}} * sizeof({{out_dtype}})
cdef {{out_dtype}} *outBufStrt
cdef {{out_dtype}} *outBufEnd
cdef {{out_dtype}} *outBuf
{%- endif %}

{% if async %}
# track start and end vars (BufVars[8] and [9]) for each tick to flush
ctypedef struct asyncBufTickInfo:
{%- for sig,dtype in sig_types.items() %}
  uint32_t {{sig}}StrtOffsetSamples  # buffered {{sig}}BufVars[0] values
  uint32_t {{sig}}StrtOffset  # buffered {{sig}}BufVars[8] values
  uint32_t {{sig}}EndOffset  # buffered {{sig}}BufVars[9] values
  uint32_t {{sig}}Len  # buffered {{sig}}Len values
{%- endfor %}

cdef uint32_t asyncStrtOffset
cdef uint32_t asyncEndOffset
cdef uint32_t *pAsyncFlushIndexStrt
cdef uint32_t *pAsyncFlushIndexEnd
cdef asyncBufTickInfo *pAsyncBufInfo
cdef uint8_t *pAsyncMem
cdef size_t asyncShmSize
ASYNC_BUF_LEN = {{async_buf_len}}
{% endif %}

# python variables
lastVal = 1
in_sigs = {}
{%- if has_in_signal == 1 %}
in_sig = None
{%- endif %}
{%- if not is_main_process %}
shouldExit = False
{%- endif %}

{%- if is_writer %}
# driver variables
{{ driver_code["variables"] }}
{%- endif %}

cdef void handle_exit(int exitStatus):
  global retVal, tid, radio, context, s, shmSize, pmem, pTickUpSem, pTickDownSem, pSigSems

{%- if is_writer %}

  {{destruct_code}}

{{ driver_code["exit_handler"] }}

{%- endif %}

  munmap(pmem, shmSize)
  munlockall()

{%- if is_main_process %}

  if (sem_close(pTickUpSem) == -1):
    printf("Could not close sink up semaphore. \n")
  if (sem_close(pTickDownSem) == -1):
    printf("Could not close sink down semaphore. \n")
  {%- for sig_sem in sig_sems %}
  if (sem_close(pSigSems[{{sig_sem[1]}}]) == -1):
    printf("Could not close sink signal semaphore. \n")
  {%- endfor %}

  shm_unlink("/async_mem_{{name}}")
{%- endif %}
  exit(exitStatus)

cdef void int_handler(int signum):
  pass

cdef void exit_handler(int signum):
  global shouldExit
{%- if is_main_process %}
  handle_exit(0)
{%- else %}
  shouldExit = True
{%- endif %}

cdef void bus_handler(int signum):
  printf("Bus Error\n")
  handle_exit(1)

cdef void segv_handler(int signum):
  printf("Seg Fault\n")
  handle_exit(1)

cdef void usr2_handler(int signum):
  pass

# main
sigemptyset(&exitMask)
sigaddset(&exitMask, SIGALRM)
init_utils(&handle_exit, &exitMask)

pid = getpid()
ppid = getppid()

set_sighandler(SIGINT, &int_handler, &exitMask)
set_sighandler(SIGUSR1, &exit_handler, &exitMask)
set_sighandler(SIGBUS, &bus_handler, &exitMask)
set_sighandler(SIGSEGV, &segv_handler, &exitMask)
set_sighandler(SIGUSR2, &usr2_handler, NULL)

# open and map shared memory
shmSize = ROUND_UP(SHM_SIZE, PAGESIZE)
open_shared_mem(&pmem, SMEM0_PATHNAME, shmSize, O_RDWR, PROT_READ | PROT_WRITE)
pNumTicks = <int64_t *>(pmem + NUM_TICKS_OFFSET)
pBufVars = <uint32_t *>(pmem + BUF_VARS_OFFSET)
{%- for sig, args in in_signals.items() %}
{{sig}}BufVars = pBufVars + {{in_sig_nums[sig] * buf_vars_len}}
{% endfor %}
{%- if is_main_process %}
pTickUpSem = sem_open("/tick_up_sem_{{non_source_num}}", 0)
pTickDownSem = sem_open("/tick_down_sem_{{non_source_num}}", 0)
  {%- for name, sem_num in in_sig_sems %}
snprintf(semNameBuf, SEM_NAME_LEN, "/sig_sem_%d", {{sem_num}})
pSigSems[{{sem_num}}] = sem_open(semNameBuf, 0)
  {%- endfor %}
{%- endif %}

{%- if async %}
# put buffer flush tracking variables in shared memory between
asyncShmSize = ROUND_UP(sizeof(asyncBufTickInfo) * ASYNC_BUF_LEN, PAGESIZE)
{% if is_main_process %}
open_shared_mem(
  &pAsyncMem,
  "/async_mem_{{name}}",
  asyncShmSize,
  O_TRUNC | O_CREAT | O_RDWR,
  PROT_READ | PROT_WRITE
)
{% else %}
open_shared_mem(
  &pAsyncMem,
  "/async_mem_{{name}}",
  asyncShmSize,
  O_RDWR,
  PROT_READ | PROT_WRITE
)
{% endif %}
pAsyncFlushIndexStrt = <uint32_t *>pAsyncMem
pAsyncFlushIndexEnd = <uint32_t *>(pAsyncMem + sizeof(uint32_t))
pAsyncBufInfo = <asyncBufTickInfo *>(pAsyncMem + 2 * sizeof(uint32_t))
pAsyncFlushIndexStrt[0] = 0
pAsyncFlushIndexEnd[0] = 0
{%- endif %}

{% if is_writer -%}

{%- if platform_system == "Linux" %}
created_mem = sa.list()
  {% for sig,args in in_signals.items() %}
if any([b'{{sig}}' == x[0] for x in created_mem]):
  in_sigs['{{sig}}'] = sa.attach("shm://{{sig}}")
  {{sig}}Raw = in_sigs['{{sig}}']
    {%- if args["max_packets_per_tick"] == 1 %}
  {{sig}} = {{sig}}Raw[{{sig}}BufVars[8]]
    {%- else %}
  {{sig}} = {{sig}}Raw[{{sig}}BufVars[8]:{{sig}}BufVars[9]]
    {%- endif %}
else:
  die("numpy signal not created\n")
    {% if out_signal["args"]["type"] != "vis_pygame" %}
{{sig}}StrtPtr = <{{args["ctype"]}} *><long>in_sigs["{{sig}}"].__array_interface__["data"][0]
{{sig}}EndPtr = {{sig}}StrtPtr + len(in_sigs["{{sig}}"])
    {% endif %}
    {% if in_signals|length == 1 %}
in_sig = in_sigs["{{sig}}"]
    {%- endif %}
  {% endfor %}
{%- else %}
  {% for sig,args in in_signals.items() %}
in_sigs["{{sig}}"] = sa.attach("shm://{{sig}}")
{{sig}}Raw = in_sigs["{{sig}}"]
    {%- if args["max_packets_per_tick"] == 1 %}
{{sig}} = {{sig}}Raw[{{sig}}BufVars[8]]
    {%- else %}
{{sig}} = {{sig}}Raw[{{sig}}BufVars[8]:{{sig}}BufVars[9]]
    {%- endif %}
    {% if out_signal["args"]["type"] != "vis_pygame" %}
{{sig}}StrtPtr = <{{args["ctype"]}} *><long>in_sigs["{{sig}}"].__array_interface__["data"][0]
{{sig}}EndPtr = {{sig}}StrtPtr + len(in_sigs["{{sig}}"])
    {% endif %}
    {% if in_signals|length == 1 %}
in_sig = in_sigs["{{sig}}"]
    {%- endif %}
  {% endfor %}
{%- endif %}

{{ driver_code["setup"] }}

{%- endif %}

{% if is_writer %}
  {%- if has_in_signal %}
outBufStrt = <{{in_signal_type}} *><long>in_sig.__array_interface__["data"][0]
outBuf = outBufStrt
outBufEnd = outBufStrt + {{in_signal_name}}BufVars[7]
  {%- else %}
outBufStrt = <{{out_dtype}} *>malloc(packetSize)
outBuf = outBufStrt
outBufEnd = outBufStrt + packetSize
  {%- endif %}
{%- endif %}

{% if default_params -%}
{{default_sig_name}} = {{default_params}}
{%- endif %}

{% if is_writer -%}
{{construct_code}}
{%- endif %}


{%- if platform_system != "Darwin" %}
make_realtime()
{%- endif %}

kill(ppid,SIGUSR1)

while(True):
{%- if is_main_process %}
  sem_wait(pTickUpSem)
  {%- for name, sem_num in in_sig_sems %}
  sem_wait(pSigSems[{{sem_num}}])
  {%- endfor %}
  {% if async %}
  # async main process updates buffer flush end variable
  if ((pAsyncFlushIndexEnd[0] + 1) == pAsyncFlushIndexStrt[0]) or (
    (pAsyncFlushIndexEnd[0] == ASYNC_BUF_LEN) and
    (pAsyncFlushIndexStrt[0] == 0)
  ):
    die("Async buffer overrun. Consider increasing signal history sizes.")
  {% for sig,args in in_signals.items() %}
  pAsyncBufInfo[pAsyncFlushIndexEnd[0]].{{sig}}StrtOffsetSamples = {{sig}}BufVars[0]
  pAsyncBufInfo[pAsyncFlushIndexEnd[0]].{{sig}}StrtOffset = {{sig}}BufVars[8]
  pAsyncBufInfo[pAsyncFlushIndexEnd[0]].{{sig}}EndOffset ={{sig}}BufVars[9]
  pAsyncBufInfo[pAsyncFlushIndexEnd[0]].{{sig}}Len = {{sig}}BufVars[1] - {{sig}}BufVars[0]
  {% endfor %}
  pAsyncFlushIndexEnd[0] += 1
  if pAsyncFlushIndexEnd[0] >= ASYNC_BUF_LEN:
    pAsyncFlushIndexEnd[0] = 0
  {% endif %}
{%- endif %}

{%- if is_writer %}

  {% if async -%} # async buffering and parser logic

  # async writer udpates buffer flush start variable
  if pAsyncFlushIndexStrt[0] == pAsyncFlushIndexEnd[0]:
    if (shouldExit) and (pAsyncFlushIndexStrt[0] == pAsyncFlushIndexEnd[0]):
      handle_exit(0)
    # wait for half a tick if there is no data
    time.sleep({{config["config"]["tick_len"]}} / (2. * 1e6))
    continue

    {% for sig, args in in_signals.items() %}
  asyncStrtOffset = pAsyncBufInfo[pAsyncFlushIndexStrt[0]].{{sig}}StrtOffset
  asyncEndOffset = pAsyncBufInfo[pAsyncFlushIndexStrt[0]].{{sig}}EndOffset
  if asyncEndOffset - asyncStrtOffset == 1:
    {{sig}} = {{sig}}Raw[asyncStrtOffset]
  else:
    {{sig}} = {{sig}}Raw[asyncStrtOffset:asyncEndOffset]
  {{sig}}Len =  pAsyncBufInfo[pAsyncFlushIndexStrt[0]].{{sig}}Len
    {% endfor %}

    {%- if has_parser %}

  {{parser_code}}
    {%- endif %}

    {%- if has_in_signal %}

  outBuf = outBufStrt + pAsyncBufInfo[pAsyncFlushIndexStrt[0]].{{in_signal_name}}StrtOffsetSamples
    {% endif %}
  pAsyncFlushIndexStrt[0] += 1
  if pAsyncFlushIndexStrt[0] >= ASYNC_BUF_LEN:
    pAsyncFlushIndexStrt[0] = 0

  {% else -%} # synchronous parser logic

  # perform parsing for each packet
    {% for sig, args in in_signals.items() %}
  if {{sig}}BufVars[9] - {{sig}}BufVars[8] == 1:
    {{sig}} = {{sig}}Raw[{{sig}}BufVars[8]]
  else:
    {{sig}} = {{sig}}Raw[{{sig}}BufVars[8]:{{sig}}BufVars[9]]
  {{sig}}Len = {{sig}}BufVars[1] - {{sig}}BufVars[0]
    {% endfor %}

    {%- if has_parser %}

  {{parser_code}}
    {%- endif %}

    {%- if has_in_signal %}

  outBuf = outBufStrt + {{in_signal_name}}BufVars[0]
    {%- endif %}

  {% endif %}


{{ driver_code["write"] }}

  {%- if not is_main_process %}
  if (shouldExit) and (pAsyncFlushIndexStrt[0] == pAsyncFlushIndexEnd[0]):
    handle_exit(0)
  {%- endif %}
{%- endif %}

{%- if is_main_process %}
  sem_post(pTickDownSem)
{%- endif %}
