# *****************************************************************************
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Module authors:
#   Markus Zolliker <markus.zolliker@psi.ch>
# *****************************************************************************
"""Module to handle running SEC nodes for an instrument

Also start and stop frappy servers related to SEA

Handled by linsetools:

    For all detected SEC nodes, setup files are written to the
    setup directory, and they are disabled ('lowlevel') when shut down.

    For SEC nodes started on the instrument computer, wrapper config files
    containing the port number are produced in the wrapper directory.
"""
import sys
sys.path.append('/home/l_samenv')
sys.path.append('/sq_sw/linse')
import time
import socket
import json
import os
from select import select
from frappy.core import Readable, Parameter, Property, Command, Attached, IDLE
from frappy.datatypes import ArrayOf, StructOf, StringType, IntRange, TupleOf, BoolType
from frappy.client import SecopClient
import linsetools.frappy
from sehistory.normalizeuri import normalizeuri


porttype = IntRange(0, 0xc000)
secnodetype = StructOf(cfg=StringType(), uri=StringType(), service=StringType())


SECOP_UDP_PORT = 10767


class FrappyControl(linsetools.frappy.FrappyControl):
    # TODO; move to linsetools
    def service_from_uri(self, uri):
        try:
            idx = self.frappy_servers.index(uri)
            return {0: 'main', 1: 'stick'}.get(idx, 'addons')
        except ValueError:
            return ''


class Listener:
    socket = None

    def __init__(self, use_localhost=False):
        self.use_localhost = use_localhost  # whether 'localhost' or the real hostname is returned on the own machine

    def poll(self, log=None):
        if self.socket is None:
            return None
        if not select([self.socket], [], [], 0)[0]:
            return None
        try:
            msg, addr = self.socket.recvfrom(1024)
        except socket.error:  # pragma: no cover
            return None
        addr = socket.getnameinfo(addr, socket.NI_NOFQDN)[0]
        msg = json.loads(msg.decode('utf-8'))
        if log:
            log.debug('got msg %r', msg)
        kind = msg.pop('SECoP', None)
        if kind == 'node':
            msg['device'] = msg['equipment_id'].split('.')[0]
            uri = f"{addr}:{msg['port']}"
        elif kind == 'for_other_node':
            uri = msg['uri']
        else:
            return None
        host, _, port = uri.rpartition(':')
        host = normalizeuri(host or 'localhost', self.use_localhost)
        msg['uri'] = f'{host}:{port}'
        return msg


class UdpScan(Listener):
    def start(self, log=None):
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        # send a general broadcast
        try:
            sock.sendto(json.dumps(dict(SECoP='discover')).encode('utf-8'),
                        ('255.255.255.255', SECOP_UDP_PORT))
        except OSError as e:
            if log:
                log.info('could not send the broadcast %r:', e)
        self.socket = sock
        self.deadline = time.time() + 30

    def poll(self, log=None):
        if self.socket is None:
            return None
        if time.time() > self.deadline:
            try:
                self.socket.close()
            except Exception:
                pass
            self.socket = None
        return super().poll(log)


class UdpListener(Listener):
    def start(self):
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.settimeout(1)
        if os.name == 'nt':
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        else:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        sock.bind(('0.0.0.0', SECOP_UDP_PORT))
        self.socket = sock


def send_other_udp(uri, instrument, device=None):
    """inform the feeder about the start of a frappy server"""
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    msg = {
        'SECoP': 'for_other_node',
        'uri': uri,
        'instrument': instrument,
    }
    if device:
        msg['device'] = device
    msg = json.dumps(msg, ensure_ascii=False, separators=(',', ':')).encode('utf-8')
    sock.sendto(msg, ('255.255.255.255', SECOP_UDP_PORT))


class SecNode:
    log = None

    def __init__(self, host_port, cfg, service=''):
        self.host_port = host_port
        self.host, _, self.port = host_port.partition(':')
        self.client = SecopClient(host_port)
        self.client.register_callback(None, self.nodeStateChange)
        self.nodename = cfg
        self.cfg = cfg
        self.description = cfg
        self.online = False
        self.status = 'created'
        self.trigger = True
        # self.announce_pnp = True  # None: do not announce, True: announce new, False: announce removal
        self.setup_was_loaded = False
        self.service = service

    def connect(self, complete_callback=None, log=None):
        self.status = 'connecting'
        self.log = log
        self.complete_callback = complete_callback
        if self.log:
            self.log.info('spawn connect')
        self.client.spawn_connect(self.complete)

    def complete(self):
        try:
            self.online = True
            self.status = 'completing'
            self.nodename = self.client.nodename
            if self.log:
                self.log.info('connected to %r', self.nodename)
            if self.complete_callback:
                if not self.cfg:
                    if self.nodename == self.client.uri:
                        self.cfg = self.host_port
                    else:
                        self.cfg = self.nodename.replace('.', '_')
                try:
                    self.complete_callback(self)
                except Exception as e:
                    self.log.exception('complete_callback failed')
                self.complete_callback = None
            desc = self.client.properties.get('description') or self.nodename
            self.description = desc.split('\n')[0]
            self.status = 'connected'
        except Exception as e:
            self.log.exception('connect failed')
            self.status = f'disconnected {e!r}'

    def nodeStateChange(self, online, statetext):
        pass

    def disconnect(self):
        self.online = False
        self.status = 'disconnecting'
        self.client.disconnect()
        self.status = 'disconnected'


class SuperFrappy(Readable):
    seascan = Attached()
    marcheport = Property('marche port number', porttype, default=8124)
    is_main_instrument = Property('this is the main instrument', BoolType(), default=True)
    value = Parameter('running servers', ArrayOf(secnodetype), default=())
    instance = Parameter('"this" or <instrument>', StringType(), default='this')
    # plugplay = Parameter('enable plug and play', BoolType(), readonly=False, default=False)
    # nicos_setups = Parameter('active nicos se setups', ArrayOf(StringType()))
    _secnodes = None  # dict <host_post> of SecNode
    _udp_listener = None
    _setups = ()
    _to_close = ()
    _fast_deadline = 0

    def initModule(self):
        super().initModule()
        # fc = FrappyControl(self.instance, 'localhost', self.marcheport)
        self._secnodes = {}
        self._udp_listener = [UdpScan(True), UdpListener(True)]
        self.rescan()
        self._to_close = set()
        for pname in 'device_name', 'stick_name', 'addons':
            self.seascan.addCallback(pname, self.sea_update, pname)

    def add_secnode(self, service, cfg, start):
        self.log.info('add %r as %r', cfg, service)
        # TODO: treat unknown service
        for secnode in self._secnodes.values():
            if secnode.service == service:
                if cfg == secnode.cfg:
                    return
            elif cfg == secnode.cfg:
                raise ValueError(f'{cfg} is already used as {secnode.service!r}')
        fc = FrappyControl(self.instance, 'localhost', self.marcheport)
        uri = fc.add_frappy_service(service, cfg, None, self.log)
        if start:
            self.log.info('start %r', cfg)
            fc.start(cfg)
        self._secnodes[uri] = secnode = SecNode(uri, cfg, service)
        secnode.connect()
        self.read_value()
        self._fast_deadline = time.time() + 15
        self.setFastPoll(True, 0.25)

    def remove_secnode(self, service_or_cfg, remaining=()):
        fc = FrappyControl(self.instance, 'localhost', self.marcheport)
        count = 0
        for uri, secnode in list(self._secnodes.items()):
            if secnode.service == service_or_cfg:
                cfg = secnode.cfg
            elif secnode.cfg == service_or_cfg:
                cfg = service_or_cfg
            else:
                continue
            if cfg not in remaining:
                count += 1
                self.log.info('remove %s at %s', secnode.cfg, uri)
                secnode.disconnect()
                fc.stop(cfg)
                fc.delete_frappy_service(cfg)
        self.read_value()
        return count

    def sea_update(self, pname, value, error):
        if error or not value:
            self.log.warn('SEA ignore %r %r %r', pname, value, error)
            return
        self.log.warn('SEA %r %r %r', pname, value, error)
        fc = FrappyControl(self.instance, 'localhost', self.marcheport)
        if pname == 'device_name':
            service = 'main'
            configs = [f'{value}.config']
        elif pname == 'stick_name':
            service = 'stick'
            configs = [f'{value}.stick']
        else:
            service = 'addons'
            configs = [f'{v}.config' for v in value.split()]
        self.log.warn('SeA %r', configs)
        if value:
            fc.all_cfg(self.instance, service, details=True)
            for seacfg in configs:
                cfgs = list(fc.sea2frappy.get(seacfg, ()))
                if len(cfgs) == 1:
                    cfg = cfgs[0]
                    self.log.warn('SeA cfg=%s', cfg)
                    try:
                        self.add_secnode(service, cfg, False)  # TODO: must be True, finally
                    except ValueError as e:
                        self.log.warn('%s', e)
                elif len(cfgs) > 1:
                    self.log.warn('frappy config for %s is ambiguous: %s', seacfg, ','.join(cfgs))
            if service == 'addons':
                remaining = {fc.sea2frappy.get(f'{v["cfg"]}.addons') for v in self.value}
            else:
                return
        else:
            remaining = ()
        self.remove_secnode(service, remaining=remaining)

    def doPoll(self):
        super().doPoll()
        done = False
        for listener in self._udp_listener:
            while msg := listener.poll(self.log):
                done = True
                uri = msg['uri']
                if uri in self._secnodes:
                    self.log.info('%r is already known', msg)
                else:
                    cfg = msg.get('device', '')
                    # if uri.startswith('localhost:') and cfg != 'superfrappy':
                    if cfg != 'superfrappy':
                        self.log.warn('%r appeared', msg)
                        try:
                            self.connect(uri)
                        except Exception as e:
                            self.log.warning('%s: %r', uri, e)
                    else:
                        self.log.info('skip %r', msg)

        if self._fast_deadline == 0:
            self.log.warn('INIT')
            try:
                for pname in 'device_name', 'stick_name', 'addons':
                    method = f'read_{pname}'
                    value = getattr(self.seascan, method)()
                    self.sea_update(pname, value, None)
                self._fast_deadline = time.time() + 15
                self.setFastPoll(True, 0.25)
            except Exception as e:
                self.log.error('doPoll %r', e)
        if not done and time.time() > self._fast_deadline:
            self.setFastPoll(False)

    def read_value(self):
        value = []
        for secnode in self._secnodes.values():
            info = dict(cfg=secnode.cfg, uri=secnode.host_port, service=secnode.service)
            value.append(info)
        if tuple(value) != self.value:
            self.log.info('----')
            for info in value:
                formatted = '%(cfg)s uri=%(uri)s service=%(service)s' % info
                self.log.info('---- %s', formatted)
        return value

    def read_status(self):
        return IDLE, ', '.join((f'{k}: {v.status}' for k, v in self._secnodes.items()))

    @Command
    def rescan(self):
        """rescan secop servers in subnet"""
        for listener in self._udp_listener:
            listener.start()

    @Command(argument=StringType())
    def restart(self, cfg):
        """restart the frappy server

        works only when it runs on the local machine
        """
        fc = FrappyControl(self.instance, 'localhost', self.marcheport)
        fc.restart(cfg)

    @Command(argument=StringType())
    def connect(self, host_port):
        """connect secnode

        :param host_port: box address <host>:<port>
        """
        service, _, host_port = host_port.rpartition('=')
        host_port = normalizeuri(host_port, True)
        secnode = self._secnodes.get(host_port)
        if secnode:
            self.log.info('already connected %r', host_port)
        else:
            self.log.info('connect secnode %r', host_port)
            fc = FrappyControl(self.instance, 'localhost', self.marcheport)
            if not service:
                service = fc.service_from_uri(host_port)
                if not service:
                    raise ValueError(f'can not determine service from {host_port}')
            secnode = SecNode(host_port, fc.cfg_info.get(host_port, ''), service)
            self._secnodes[host_port] = secnode
            secnode.connect(log=self.log)
        self.read_value()

    @Command(argument=StringType())
    def main(self, cfg):
        """add or change main cfg"""
        self.add_secnode('main', cfg, True)

    @Command(argument=StringType())
    def stick(self, cfg):
        """add or change stick cfg"""
        self.add_secnode('stick', cfg, True)

    @Command(argument=StringType())
    def add(self, cfg):
        """add an addons cfg"""
        self.add_secnode('addons', cfg, True)

    @Command(argument=StringType())
    def remove(self, service_or_cfg):
        """remove server

        - for servers on localhost stop the server
        - register the setup file for deletion
        """
        count = self.remove_secnode(service_or_cfg)
        if count == 0:
            raise ValueError(f'no running server found for {service_or_cfg}')

    @Command(argument=TupleOf(StringType(), porttype), result=StringType())
    def running_cfg(self, host_port):
        """when running, return cfg, else an empty string"""
        secnode = self._secnodes.get(host_port)
        if secnode and secnode.online:
            return secnode.cfg
        return ''
