# *****************************************************************************
# 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>
#   Anik Stark <anik.stark@psi.ch>
# *****************************************************************************

import struct
from frappy.core import BytesIO, HasIO, Drivable, Parameter, Property, IntRange, FloatRange, BoolType, \
    Attached, IDLE, BUSY, ERROR, Command, TupleOf, StringType
from frappy.errors import CommunicationFailedError


class Item:

    def __init__(self, unit, description, reference, value):
        self.unit = unit
        self.description = description
        self.reference = reference
        self.value = value


class IO(BytesIO):

    db = Property('database number', datatype=IntRange(0, 255), default=200)
    _db_length = 0
    fetch_header = struct.pack('>2s14B', b'S5', 16, 1, 3, 6, 15, 3, 0, 255, 7, 0, 0, 0, 0, 0)
    data = None

    def checkHWIdent(self):
        self.data = {}
        reply = self.send_fetch_message(0, 1)    # length in words
        self._db_length = struct.unpack('>H', reply)[0] // 2    # length in words (1 word = 2 bytes = 16 bit)
        self.get_data()

    def send_fetch_message(self, start, length):
        """fetch data from start (byte offset) with length (word length = byte-length / 2)"""
        msg = struct.pack('>2s8B2H2B', b'S5', 16, 1, 3, 5, 3, 8, 1, self.db, start, length, 255, 2)
        reply = self.communicate(msg, 16 + length * 2)
        if reply[:16] != self.fetch_header:
            raise CommunicationFailedError('bad reply header: %r' % reply[:16])
        return reply[16:]

    @Command(argument=TupleOf(IntRange(0,999), IntRange(0,255), IntRange(0, 255)),
             result=StringType())
    def write(self, start, byte1, byte2):
        """write 2 bytes, for debug purposes"""
        msg = struct.pack('>2s8BHH4B', b'S5', 16, 1, 3, 3, 3, 8, 1, self.db, start, 1, 255, 2, byte1, byte2)
        return repr(self.communicate(msg, 16))

    @Command(argument=IntRange(0, 999), result=StringType())  # result = TupleOf(*[IntRange(0, 255) for _ in range(4)])
    def read(self, start):
        """"read 16 bytes, for debug purposes"""
        msg = struct.pack('>2s8BHH2B', b'S5', 16, 1, 3, 5, 3, 8, 1, self.db, start, 8, 255, 2)
        reply = self.communicate(msg, 16+16)
        return ' '.join(f'{v:02x}' for v in reply[16:])

    def send_write_message(self, key, value):
        info = self.data[key]
        name = info['name'].encode('latin-1')
        length = len(name)
        # check if name matches
        reply = self.send_fetch_message(info['name_pos'], length // 2 + 2)
        if reply[1] != length and reply[2:2+length] != name:
            raise CommunicationFailedError('communication error, name mismatch')
        typ = info['typ']
        if typ == 1:    # int16 (or uint16 ?)
            content = struct.pack('>h', value)
        if typ == 4:    # bool
            content = struct.pack('>H', bool(value) * 256)
        elif typ == 2:   # int32 (or uint32 ?)
            content = struct.pack('>i', value)
        elif typ == 3:   # float
            content = struct.pack('>f', value)
        else:
            raise ValueError('unknown type')
        self.log.debug('write %s: %r %r', key, info['data_pos'], content)
        msg = struct.pack('>2s8BHH2B', b'S5', 16, 1, 3, 3, 3, 8, 1, self.db,
                          info['data_pos'], info['data_size'], 255, 2)
        reply = self.send_fetch_message(info['data_pos']-16, 16)
        self.log.debug('old %r %r', reply[:16], reply[16:])
        self.communicate(msg + content, 16)
        reply = self.send_fetch_message(info['data_pos']-16, 16)
        self.log.debug('new %r %r', reply[:16], reply[16:])

    def get_string(self, pos, reply):
        size, length = reply[pos:pos+2]    # reserved space, actual length (containing info)
        pos += 2
        value = reply[pos:pos+length].decode('latin-1')
        pos += size
        return pos, value

    def get_data(self):
        reply = self.send_fetch_message(0, self._db_length)
        # self.log.info('get_data reply %s', reply)
        pos = 2
        result = {}
        while pos < len(reply):
            typ, alarm = reply[pos:pos+2]
            pos += 2
            name_pos = pos
            pos, name = self.get_string(pos, reply)
            pos, unit = self.get_string(pos, reply)
            pos, description = self.get_string(pos, reply)
            pos, reference = self.get_string(pos, reply)
            data_pos = pos
            if typ == 1:  # int16 (or uint16?)
                length = 2
                value = struct.unpack('>h', reply[pos:pos+length])[0]
            elif typ == 4:  # bool
                length = 2
                value = bool(struct.unpack('>H', reply[pos:pos+length])[0])
            elif typ == 2:  # int32 (or uint32?)
                length = 4
                value = struct.unpack('>i', reply[pos:pos+length])[0]
            elif typ == 3:  # float32
                data_pos = pos + 2  # skip 2 bytes with unknown purpose
                length = 4
                value = struct.unpack('>f', reply[data_pos:data_pos+length])[0]
            else:
                raise ValueError('unknown type')
            pos = data_pos + length
            res = {
                'name': name,
                'unit': unit,
                'description': description,
                'reference': reference,
                'typ': typ,
                'alarm': alarm,
                'value': value,
                'name_pos': name_pos,
                'data_pos': data_pos,
                'data_size': length // 2,
            }
            if res != self.data.get((name, description)):
                self.log.debug('%r, %r', reply[data_pos:pos], res)
            result[name, description] = res
        self.data = result


class Magnet(HasIO, Drivable):

    """"
    Write-message for target (concatenated from the following snippets)
        sequence                            size    meaning
        --------                            ----    -------
        b'S5\x10\x01\x03\x03\x03\x08\x01'   9       fixed protocol header
        b'\xc8'                             1       db number ?
        b'\x00\x46'                         2       start in bytes, i.e. 70 for setpoint
        b'\x00\x04'                         2       length in words (= bytes / 2), i.e. 4 for setpoint
        b'\xff\x02'                         2       fixed
        b'\x03\x00\x0e\x0c'                 var.    payload data, size according to type?
    """

    ioClass = IO

    target = Parameter('target', datatype=FloatRange(unit='T'))
    value = Parameter('value', datatype=FloatRange(unit='T'))
    scale = Parameter('scale for actual value', datatype=FloatRange(unit='T/A'), default=0.3/75)
    ramp_time = Parameter('time to ramp to 75A', datatype=FloatRange(unit='sec'))  # seems to be readonly
    main_switch = Parameter('main switch', datatype=BoolType())  # seems to be readonly
    emergency_stop = Parameter('state of emergency stop button', datatype=BoolType())
    collective_fault = Parameter('collecitve fault flag', datatype=BoolType())
    holder_1_inserted = Parameter('sample holder 1 is inserted', datatype=BoolType())
    holder_2_inserted = Parameter('sample holder 2 is inserted', datatype=BoolType())
    motor_stop = Parameter('sample holder motor stop flag', datatype=BoolType())
    current_timeout = Parameter('magnet current timeout flag', datatype=BoolType())
    polarity = Parameter('positive polarity', datatype=BoolType())
    _busy = False

    def extract_data(self, key):
        val = self.io.data[key]['value']
        return float(val)

    def read_value(self):
        current = self.extract_data(('actual_value', 'Magnetcurrent'))
        return current * self.scale

    def read_target(self):
        return self.extract_data(('setpoint', 'Magnetcurrent')) * self.scale

    def write_target(self, target):
        self._busy = True
        self.status = BUSY, 'changed target'
        self.setFastPoll(True)
        # self.write_main_switch(True)
        self.io.send_write_message(('setpoint', 'Magnetcurrent'), target / self.scale)

    def read_ramp_time(self):
        return self.extract_data(('time', 'Current_ramp@75A'))

    # def write_ramp_time(self, ramp_time):
    #     self.io.send_write_message(('time', 'Current_ramp@75A'), ramp_time)

    def read_collective_fault(self):
        return self.extract_data(('alert', 'collective_fault'))

    def read_emergency_stop(self):
        return self.extract_data(('alert', 'emergency_stop'))

    # def write_main_switch(self, val):
    #     self.io.send_write_message(('main_switch', 'on'), val)

    def read_main_switch(self):
        return self.extract_data(('main_switch', 'on'))

    def read_holder_1_inserted(self):
        return self.extract_data(('sample_holder', 'Type1_in_position'))

    def read_holder_2_inserted(self):
        return self.extract_data(('sample_holder', 'Type2_in_position'))

    def read_motor_stop(self):
        return self.extract_data(('sample_holder', 'Motor_stop'))

    def read_current_timeout(self):
        return self.extract_data(('magnet_current', 'Timeout'))

    def read_polarity(self):
        return self.extract_data(('magnet_current', 'Polarity_pos'))

    def read_status(self):
        if self.collective_fault:
            self._busy = False
            self.setFastPoll(False)
            return ERROR, 'collective fault'
        if self.emergency_stop:
            self._busy = False
            self.setFastPoll(False)
            return ERROR, 'emergency stop is active'
        if self._busy:
            if abs(self.value - self.target) < 0.1:
                self._busy = False
                self.setFastPoll(False)
            return BUSY, 'ramping'
        return IDLE, ''

    def stop(self):
        """stop ramp"""
        self.target = self.read_value()
        self.io.send_write_message(('setpoint', 'Magnetcurrent'), self.target)

    def doPoll(self):
        self.io.get_data()
        super().doPoll()


class SampleHolder(Drivable):

    motor = Attached()

    target = Parameter('target sample holder position', datatype=IntRange(0,9), readonly=False)
    value = Parameter('sample holder position', datatype=IntRange(0,9))
    target_max = Parameter('highest sample holder number', datatype=IntRange(0,9), default=6, readonly=False)
    pos_offset = Parameter('position of first sample holder [mm]', datatype=FloatRange(unit='mm'), default=43)
    pos_diff = Parameter('distance between two sample positions [mm]', datatype=FloatRange(unit='mm'), default=50)
    tolerance = Parameter('position tolerance', datatype=FloatRange(unit='mm'), default=1)
    _pos = None     # position of sample holder in units of target (where 1 unit is eq. to pos_diff)
    _stopflag = False

    def read_value(self):
        if not self.motor.isBusy():
            val = self.motor.value
            self._pos = (val - self.pos_offset) / self.pos_diff    # convert mm to sample position
            return round(self._pos)
        return self.value

    def write_target(self, target):
        self._stopflag = False
        val = self.pos_offset + target * self.pos_diff
        self.motor.write_target(val)
        self.read_status()
        self.setFastPoll(True)

    def read_status(self):
        if self.motor.isBusy():
            return BUSY, ('stopping motor' if self._stopflag else 'motor is running')
        self.setFastPoll(False)
        if abs(self._pos - self.value) > abs(self.tolerance / self.pos_diff):
            if self._stopflag:
                return ERROR, 'stopped at intermediate position'
            return ERROR, 'intermediate position'
        return IDLE, ''

    def stop(self):
        self.motor.stop()
        self._stopflag = True
        self.status = BUSY, 'stopping motor'
