#!/usr/bin/env python3

# This file is part of OpenModelica.
#
# Copyright (c) 1998-2026, Open Source Modelica Consortium (OSMC),
# c/o Linköpings universitet, Department of Computer and Information Science,
# SE-58183 Linköping, Sweden.
#
# All rights reserved.
#
# THIS PROGRAM IS PROVIDED UNDER THE TERMS OF AGPL VERSION 3 LICENSE OR
# THIS OSMC PUBLIC LICENSE (OSMC-PL) VERSION 1.8.
# ANY USE, REPRODUCTION OR DISTRIBUTION OF THIS PROGRAM CONSTITUTES
# RECIPIENT'S ACCEPTANCE OF THE OSMC PUBLIC LICENSE OR THE GNU AGPL
# VERSION 3, ACCORDING TO RECIPIENTS CHOICE.
#
# The OpenModelica software and the OSMC (Open Source Modelica Consortium)
# Public License (OSMC-PL) are obtained from OSMC, either from the above
# address, from the URLs:
# http://www.openmodelica.org or
# https://github.com/OpenModelica/ or
# http://www.ida.liu.se/projects/OpenModelica,
# and in the OpenModelica distribution.
#
# GNU AGPL version 3 is obtained from:
# https://www.gnu.org/licenses/licenses.html#GPL
#
# This program is distributed WITHOUT ANY WARRANTY; without
# even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE, EXCEPT AS EXPRESSLY SET FORTH
# IN THE BY RECIPIENT SELECTED SUBSIDIARY LICENSE CONDITIONS OF OSMC-PL.
#
# See the full OSMC Public License conditions for more details.

'''OMSimulator Server for interactive simulation of SSP files.'''

import argparse
import json
import math
import threading
import time
import sys

from OMSimulator import SSP, CRef, Capi
import zmq

__author__ = 'Lennart Ochel <lennart.ochel@ri.se>'
__version__ = Capi.getVersion()
__copyright__ = '''\
Copyright (c) 2018-CurrentYear, Open Source Modelica Consortium (OSMC),
c/o Linköpings universitet, Department of Computer and Information Science,
SE-58183 Linköping, Sweden.'''
__status__ = 'development'

def close_socket_gracefully(socket):
  '''graceful termination'''
  socket.setsockopt(zmq.LINGER, 0) # to avoid hanging infinitely
  socket.close()

def mogrify(topic: str, msg: dict):
  '''Combines a topic identifier and a json representation of a dictionary'''
  return f'{topic} {json.dumps(msg)}'

class Server:
  def __init__(self, model, result_file, interactive, endpoint_pub, endpoint_rep):
    self.print(f'OMS Server {__version__}')
    self.print(f'ZMQ {zmq.zmq_version()}')

    self._alive = True
    self._context = zmq.Context()
    self._model = SSP(model)
    self._mutex = threading.Lock()
    self._pause = interactive
    self._socket_rep = None
    self._socket_pub = None
    self._thread = None
    self._activeSignals = set()

    if result_file:
      self._model.resultFile = result_file

    #self.print(f'Model loaded: {self._model.list()}')
    # extract all available signals
    #self._signals = self._model.getAllSignals()
    #self._signals = [CRef('test', 'step', 'height'), CRef('test', 'step', 'offset'), CRef('test', 'step', 'y')]
    self._signals = []
    if interactive and not endpoint_rep:
      self.print('flag --endpoint-rep is mandatory in interactive simulation mode')
      sys.exit(1)

    # connect the PUB socket
    if endpoint_pub:
      self._socket_pub = self._context.socket(zmq.PUB)  #pylint: disable=no-member
      self._socket_pub.connect(endpoint_pub)
      self.print('PUB socket connected to {}'.format(endpoint_pub))

    # connet the REP socket
    if endpoint_rep:
      self._socket_rep = self._context.socket(zmq.REP)  #pylint: disable=no-member
      self._socket_rep.connect(endpoint_rep)
      self._socket_rep.RCVTIMEO = 1000  #in milliseconds
      self.print('REP socket connected to {}'.format(endpoint_rep))
      self._thread = threading.Thread(target=self._main, daemon=True)
      self._thread.start()

  def print(self, msg):
    print('server:  {}'.format(msg), flush=True)

  def pub_msg(self, topic, msg: dict):
    if self._socket_pub:
      msg_ = mogrify(topic, msg)
      try:
        self._socket_pub.send_string(msg_)
      except zmq.error.ZMQError as error:
        self.print('Error publishing: ' + str(error))

  def _main(self):
    alive = True
    while alive:
      try:
        msg = self._socket_rep.recv_json()
      except zmq.error.Again as error:
        #self.print('recv: ' + str(error))
        continue

      fcn = msg['fcn'] if 'fcn' in msg else ''
      answer = {'status': 'nack', 'error': 'unknown'}

      if 'simulation' == fcn:
        arg = msg['arg']
        if 'pause' == arg:
          with self._mutex:
            self._pause = True
          answer = {'status': 'ack'}
        elif 'continue' == arg:
          with self._mutex:
            self._pause = False
          answer = {'status': 'ack'}
        elif 'end' == arg:
          alive = False
          with self._mutex:
            self._alive = False
          answer = {'status': 'ack'}
      elif 'signals' == fcn:
        arg = msg['arg']
        if 'available' == arg:
          answer = {'status': 'ack', 'result': self._signals}
        elif 'enable' == arg:
          self._activeSignals.add(msg['cref'])
          answer = {'status': 'ack'}
        elif 'disable' == arg:
          self._activeSignals.remove(msg['cref'])
          answer = {'status': 'ack'}

      try:
        self._socket_rep.send_json(answer)
      except zmq.error.ZMQError as error:
        self.print('send: ' + str(error))
      else:
        self.print(msg)

  def run(self):
    self.pub_msg('status', {'progress': 0})

    inst_model = self._model.instantiate()
    inst_model.setResultFile(self._model.activeVariant.resultFile)

    time_ = float(inst_model.getTime())
    startTime = float(self._model.activeVariant.startTime)
    stopTime = float(self._model.activeVariant.stopTime)

    inst_model.setStartTime(startTime)
    inst_model.setStopTime(stopTime)

    inst_model.initialize()

    # for signal in self._signals:
    #   print("enable signal", signal, flush=True)
    #   print(f"Value for signal {signal}: {inst_model.getValue(signal)}")

    while True:
      if self._pause:
        time.sleep(0.3)
      else:
        try:
          progress = math.floor((time_-startTime) / (stopTime-startTime) * 100)
        except ZeroDivisionError:
          progress = 100

        self.pub_msg('status', {'progress': progress})
        if self._activeSignals:
          results = {'time': time_}
          for signal in self._activeSignals:
            # type_ = self._signals[signal]['type']
            # if type_ == 'Real':
            #   results[signal] = oms.getReal(signal)
            # elif type_ == 'Integer':
            #   results[signal] = oms.getInteger(signal)
            # elif type_ == 'Boolean':
            #   results[signal] = oms.getBoolean(signal)
            results[signal] = inst_model.getValue(signal)
          self.pub_msg('results', results)
        with self._mutex:
          inst_model.doStep()
          time_ = float(inst_model.getTime())
          #print(f"time: {time_} : {stopTime}", flush=True)
          if time_ >= stopTime and self._alive:
            break

    inst_model.terminate()
    inst_model.delete()
    self.pub_msg('status', {'progress': 100})

def _main():
  # parse command-line arguments
  parser = argparse.ArgumentParser(description='OMS-SERVER', allow_abbrev=False)
  parser.add_argument('--model', default=None, required=True, help='define the model to simulate')
  parser.add_argument('--endpoint-pub', default=None, help='define the endpoint for the pub/sub communication')
  parser.add_argument('--endpoint-rep', default=None, help='define the endpoint used for the interactive simulation communication')
  parser.add_argument('--interactive', action='store_true', help='enables the interactive simulation')
  parser.add_argument('--logLevel', default=0, type=int, help='defines the logging level')
  parser.add_argument('--option', default='', nargs='+', help='defines optional command line options')
  parser.add_argument('--result-file', default=None, help='defines the result file')
  parser.add_argument('--temp', default=None, help='defines the temp directory')
  parser.add_argument('--log-file', default=None, help='defines the log file')
  parser.add_argument('--working-directory', default=None, help='defines the working directory')
  args = parser.parse_args()

  Capi.setCommandLineOption(' '.join(list(map((lambda x: x[1:-1]), args.option))))
  Capi.setLoggingLevel(args.logLevel)

  if args.temp:
    Capi.setTempDirectory(args.temp)
  if args.log_file:
    Capi.setLogFile(args.log_file)
  if args.working_directory:
    Capi.setWorkingDirectory(args.working_directory)

  server = Server(args.model, args.result_file, args.interactive, args.endpoint_pub, args.endpoint_rep)
  # run simulation thread
  server.run()

if __name__ == '__main__':
  _main()
