
"""
    Copyright (C) 2017 SunSpec Alliance

    Permission is hereby granted, free of charge, to any person obtaining a
    copy of this software and associated documentation files (the "Software"),
    to deal in the Software without restriction, including without limitation
    the rights to use, copy, modify, merge, publish, distribute, sublicense,
    and/or sell copies of the Software, and to permit persons to whom the
    Software is furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included
    in all copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
    THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
    IN THE SOFTWARE.
"""

import os
import socket
import struct
import serial

try:
    import xml.etree.ElementTree as ET
except:
    import elementtree.ElementTree as ET

import sunspec.core.modbus.mbmap as mbmap

PARITY_NONE = 'N'
PARITY_EVEN = 'E'

REQ_COUNT_MAX = 125

FUNC_READ_HOLDING = 3
FUNC_READ_INPUT = 4
FUNC_WRITE_MULTIPLE = 16

TEST_NAME = 'test_name'

modbus_rtu_clients = {}

class ModbusClientError(Exception):
    pass

class ModbusClientTimeout(ModbusClientError):
    pass

class ModbusClientException(ModbusClientError):
    pass

def modbus_rtu_client(name=None, baudrate=None, parity=None):

    global modbus_rtu_clients

    client = modbus_rtu_clients.get(name)
    if client is not None:
        if baudrate is not None and client.baudrate != baudrate:
            raise ModbusClientError('Modbus client baudrate mismatch')
        if parity is not None and client.parity != parity:
            raise ModbusClientError('Modbus client parity mismatch')
    else:
        if baudrate is None:
            baudrate = 9600
        if parity is None:
            parity = PARITY_NONE

        client = ModbusClientRTU(name, baudrate, parity)
        modbus_rtu_clients[name] = client

    return client

def modbus_rtu_client_remove(name=None):

    global modbus_rtu_clients

    if modbus_rtu_clients.get(name):
        del modbus_rtu_clients[name]

class ModbusClientRTU(object):

    def __init__(self, name='/dev/ttyUSB0', baudrate=9600, parity=None):
        self.name = name
        self.baudrate = baudrate
        self.parity = parity
        self.serial = None
        self.timeout = .5
        self.write_timeout = .5
        self.devices = {}

        self.open()
        
    def open(self):

        try:
            if self.parity == PARITY_EVEN:
                parity = serial.PARITY_EVEN
            else:
                parity = serial.PARITY_NONE

            if self.name != TEST_NAME:
                self.serial = serial.Serial(port = self.name, baudrate=self.baudrate,
                                            bytesize=8, parity=parity,
                                            stopbits=1, xonxoff=0,
                                            timeout=self.timeout, writeTimeout=self.write_timeout)
            else:
                import sunspec.core.test.fake.serial as fake
                self.serial = fake.Serial(port = self.name, baudrate=self.baudrate,
                                          bytesize=8, parity=parity,
                                          stopbits=1, xonxoff=0,
                                          timeout=self.timeout, writeTimeout=self.write_timeout)

        except Exception, e:
            if self.serial is not None:
                self.serial.close()
                self.serial = None
            raise ModbusClientError('Serial init error: %s' % str(e))

    def close(self):

        try:
            if self.serial is not None:
                self.serial.close()
        except Exception, e:
            raise ModbusClientError('Serial close error: %s' % str(e))

    def add_device(self, slave_id, device):

        self.devices[slave_id] = device

    def remove_device(self, slave_id):

        if self.devices.get(slave_id):
            del self.devices[slave_id]

        # if no more devices using the client interface, close and remove the client
        if len(self.devices) == 0:
            self.close()
            modbus_rtu_client_remove(self.name)

    def _read(self, slave_id, addr, count, op=FUNC_READ_HOLDING, trace_func=None):
        resp = ''
        len_remaining = 5
        len_found = False
        except_code = None

        req = struct.pack('>BBHH', int(slave_id), op, int(addr), int(count))
        req += struct.pack('>H', computeCRC(req))
        
        if trace_func:
            s = '%s:%s[addr=%s] ->' % (self.name, str(slave_id), addr)
            for c in req:
                s += '%02X' % (ord(c))
            trace_func(s)

        self.serial.flushInput()
        try:
            self.serial.write(req)
        except Exception, e:
            raise ModbusClientError('Serial write error: %s' % str(e))

        while len_remaining > 0:
            c = self.serial.read(len_remaining)
            len_read = len(c);
            if len_read > 0:
                resp += c
                len_remaining -= len_read
                if len_found is False and len(resp) >= 5:
                    if not (ord(resp[1]) & 0x80):
                        len_remaining = (ord(resp[2]) + 5) - len(resp)
                        len_found = True
                    else:
                        except_code = ord(resp[2])
            else:
                raise ModbusClientTimeout('Response timeout')

        if trace_func:
            s = '%s:%s[addr=%s] <--' % (self.name, str(slave_id), addr)
            for c in resp:
                s += '%02X' % (ord(c))
            trace_func(s)

        crc = (ord(resp[-2]) << 8) | ord(resp[-1])
        if not checkCRC(resp[:-2], crc):
            raise ModbusClientError('CRC error')

        if except_code:
            raise ModbusClientException('Modbus exception %d' % (except_code))

        return resp[3:-2]

    def read(self, slave_id, addr, count, op=FUNC_READ_HOLDING, trace_func=None, max_count=REQ_COUNT_MAX):
        resp = ''
        read_count = 0
        read_offset = 0

        if self.serial is not None:
            while (count > 0):
                if count > max_count:
                    read_count = max_count
                else:
                    read_count = count
                data = self._read(slave_id, addr + read_offset, read_count, op=op, trace_func=trace_func)
                if data:
                    resp += data
                    count -= read_count
                    read_offset += read_count
                else:
                    return
        else:
            raise ModbusClientError('Client serial port not open: %s' % self.name)

        return resp

    def _write(self, slave_id, addr, data, trace_func=None):
        resp = ''
        len_remaining = 5
        len_found = False
        except_code = None
        func = FUNC_WRITE_MULTIPLE
        len_data = len(data)
        count = len_data/2

        req = struct.pack('>BBHHB', int(slave_id), func, int(addr), count, len_data)
        req += data
        req += struct.pack('>H', computeCRC(req))

        if trace_func:
            s = '%s:%s[addr=%s] ->' % (self.name, str(slave_id), addr)
            for c in req:
                s += '%02X' % (ord(c))
            trace_func(s)

        self.serial.flushInput()
        try:
            self.serial.write(req)
        except Exception, e:
            raise ModbusClientError('Serial write error: %s' % str(e))

        while len_remaining > 0:
            c = self.serial.read(len_remaining)
            len_read = len(c);
            if len_read > 0:
                resp += c
                len_remaining -= len_read
                if len_found is False and len(resp) >= 5:
                    if not (ord(resp[1]) & 0x80):
                        len_remaining = 8 - len(resp)
                        len_found = True
                    else:
                        except_code = ord(resp[2])
            else:
                raise ModbusClientTimeout('Response timeout')

        if trace_func:
            s = '%s:%s[addr=%s] <--' % (self.name, str(slave_id), addr)
            for c in resp:
                s += '%02X' % (ord(c))
            trace_func(s)


        crc = (ord(resp[-2]) << 8) | ord(resp[-1])
        if not checkCRC(resp[:-2], crc):
            raise ModbusClientError('CRC error')

        if except_code:
            raise ModbusClientException('Modbus exception: %d' % (except_code))
        else:
            resp_slave_id, resp_func, resp_addr, resp_count, resp_crc = struct.unpack('>BBHHH', resp)
            if resp_slave_id != slave_id or resp_func != func or resp_addr != addr or resp_count != count:
                raise ModbusClientError('Mobus response format error')

    def write(self, slave_id, addr, data, trace_func=None, max_count=REQ_COUNT_MAX):
        write_count = 0
        write_offset = 0
        count = len(data)/2

        if self.serial is not None:
            while (count > 0):
                if count > max_count:
                    write_count = max_count
                else:
                    write_count = count
                self._write(slave_id, addr + write_offset, data[(write_offset * 2):((write_offset + write_count) * 2)], trace_func=trace_func)
                count -= write_count
                write_offset += write_count
        else:
            raise ModbusClientError('Client serial port not open: %s' % self.name)

class ModbusClientDeviceRTU(object):

    def __init__(self, slave_id, name, baudrate=None, parity=None, timeout=None, ctx=None, trace_func=None, max_count=REQ_COUNT_MAX):
        self.slave_id = slave_id
        self.name = name
        self.client = None
        self.ctx = ctx
        self.trace_func = trace_func
        self.max_count = max_count

        self.client = modbus_rtu_client(name, baudrate, parity)
        if self.client is None:
            raise ModbusClientError('No modbus rtu client set for device')
        self.client.add_device(self.slave_id, self)

        if timeout is not None and self.client.serial is not None:
            self.client.serial.timeout = timeout
            self.client.serial.writeTimeout = timeout

    def close(self):

        if self.client:
            self.client.remove_device(self.slave_id)

    def read(self, addr, count, op=FUNC_READ_HOLDING):

        return self.client.read(self.slave_id, addr, count, op=op, trace_func=self.trace_func, max_count=self.max_count)

    def write(self, addr, data):

        return self.client.write(self.slave_id, addr, data, trace_func=self.trace_func, max_count=self.max_count)

TCP_HDR_LEN = 6
TCP_RESP_MIN_LEN = 3
TCP_HDR_O_LEN = 4
TCP_READ_REQ_LEN = 6
TCP_WRITE_MULT_REQ_LEN = 7

TCP_DEFAULT_PORT = 502
TCP_DEFAULT_TIMEOUT = 2

class ModbusClientDeviceTCP(object):

    def __init__(self, slave_id, ipaddr, ipport=502, timeout=None, ctx=None, trace_func=None, max_count=REQ_COUNT_MAX, test=False):
        self.slave_id = slave_id
        self.ipaddr = ipaddr
        self.ipport = ipport
        self.timeout = timeout
        self.ctx = ctx
        self.socket = None
        self.trace_func = trace_func
        self.max_count = max_count

        if ipport is None:
            self.ipport = TCP_DEFAULT_PORT
        if timeout is None:
            self.timeout = TCP_DEFAULT_TIMEOUT

        if test:
            import sunspec.core.test.fake.socket as fake
            self.socket = fake.socket()

    def close(self):

        self.disconnect()

    def connect(self, timeout=None):

        if self.socket:
            self.disconnect()

        if timeout is None:
            timeout = self.timeout

        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(timeout)
            self.socket.connect((self.ipaddr, self.ipport))
        except Exception, e:
            raise ModbusClientError('Connection error: %s' % str(e))

    def disconnect(self):

        try:
            if self.socket:
                self.socket.close()
            self.socket = None
        except Exception:
            pass

    def _read(self, addr, count, op=FUNC_READ_HOLDING):

        resp = ''
        len_remaining = TCP_HDR_LEN + TCP_RESP_MIN_LEN
        len_found = False
        except_code = None

        req = struct.pack('>HHHBBHH', 0, 0, TCP_READ_REQ_LEN, int(self.slave_id), op, int(addr), int(count))

        if self.trace_func:
            s = '%s:%s:%s[addr=%s] ->' % (self.ipaddr, str(self.ipport), str(self.slave_id), addr)
            for c in req:
                s += '%02X' % (ord(c))
            self.trace_func(s)

        try:
            self.socket.sendall(req)
        except Exception, e:
            raise ModbusClientError('Socket write error: %s' % str(e))

        while len_remaining > 0:
            c = self.socket.recv(len_remaining)
            # print 'c = {0}'.format(c)
            len_read = len(c);
            if len_read > 0:
                resp += c
                len_remaining -= len_read
                if len_found is False and len(resp) >= TCP_HDR_LEN + TCP_RESP_MIN_LEN:
                    data_len = struct.unpack('>H', resp[TCP_HDR_O_LEN:TCP_HDR_O_LEN + 2])
                    len_remaining = data_len[0] - (len(resp) - TCP_HDR_LEN)
            else:
                raise ModbusClientError('Response timeout')

        if not (ord(resp[TCP_HDR_LEN + 1]) & 0x80):
            len_remaining = (ord(resp[TCP_HDR_LEN + 2]) + TCP_HDR_LEN) - len(resp)
            len_found = True
        else:
            except_code = ord(resp[TCP_HDR_LEN + 2])

        if self.trace_func:
            s = '%s:%s:%s[addr=%s] <--' % (self.ipaddr, str(self.ipport), str(self.slave_id), addr)
            for c in resp:
                s += '%02X' % (ord(c))
            self.trace_func(s)

        if except_code:
            raise ModbusClientException('Modbus exception %d' % (except_code))

        return resp[(TCP_HDR_LEN + 3):]

    def read(self, addr, count, op=FUNC_READ_HOLDING):

        resp = ''
        read_count = 0
        read_offset = 0
        local_connect = False

        if self.socket is None:
            local_connect = True
            self.connect(self.timeout)

        try:
            while (count > 0):
                if count > self.max_count:
                    read_count = self.max_count
                else:
                    read_count = count
                data = self._read(addr + read_offset, read_count, op=op)
                if data:
                    resp += data
                    count -= read_count
                    read_offset += read_count
                else:
                    break
        finally:
            if local_connect:
                self.disconnect()

        return resp

    def _write(self, addr, data):

        resp = ''
        len_remaining = TCP_HDR_LEN + TCP_RESP_MIN_LEN
        len_found = False
        except_code = None
        func = FUNC_WRITE_MULTIPLE

        write_len = len(data)
        write_count = write_len/2
        req = struct.pack('>HHHBBHHB', 0, 0, TCP_WRITE_MULT_REQ_LEN + write_len, int(self.slave_id), func, int(addr), write_count, write_len)
        req += data

        if self.trace_func:
            s = '%s:%s:%s[addr=%s] ->' % (self.ipaddr, str(self.ipport), str(self.slave_id), addr)
            for c in req:
                s += '%02X' % (ord(c))
            self.trace_func(s)

        try:
            self.socket.sendall(req)
        except Exception, e:
            raise ModbusClientError('Socket write error: %s' % str(e))

        while len_remaining > 0:
            c = self.socket.recv(len_remaining)
            # print 'c = {0}'.format(c)
            len_read = len(c);
            if len_read > 0:
                resp += c
                len_remaining -= len_read
                if len_found is False and len(resp) >= TCP_HDR_LEN + TCP_RESP_MIN_LEN:
                    data_len = struct.unpack('>H', resp[TCP_HDR_O_LEN:TCP_HDR_O_LEN + 2])
                    len_remaining = data_len[0] - (len(resp) - TCP_HDR_LEN)
            else:
                raise ModbusClientTimeout('Response timeout')

        if not (ord(resp[TCP_HDR_LEN + 1]) & 0x80):
            len_remaining = (ord(resp[TCP_HDR_LEN + 2]) + TCP_HDR_LEN) - len(resp)
            len_found = True
        else:
            except_code = ord(resp[TCP_HDR_LEN + 2])

        if self.trace_func:
            s = '%s:%s:%s[addr=%s] <--' % (self.ipaddr, str(self.ipport), str(self.slave_id), addr)
            for c in resp:
                s += '%02X' % (ord(c))
            self.trace_func(s)

        if except_code:
            raise ModbusClientException('Modbus exception: %d' % (except_code))

    def write(self, addr, data):

        write_count = 0
        write_offset = 0
        local_connect = False
        count = len(data)/2

        if self.socket is None:
            local_connect = True
            self.connect(self.timeout)

        try:
            while (count > 0):
                if count > self.max_count:
                    write_count = self.max_count
                else:
                    write_count = count
                self._write(addr + write_offset, data[(write_offset * 2):((write_offset + write_count) * 2)])
                count -= write_count
                write_offset += write_count
        finally:
            if local_connect:
                self.disconnect()

class ModbusClientDeviceMapped(object):

    def __init__(self, slave_id, name, pathlist=None, max_count=None, ctx=None):

        self.slave_id = slave_id
        self.name = name
        self.ctx = ctx
        self.modbus_map = None

        if name is not None:
            self.modbus_map = mbmap.ModbusMap(slave_id)
            self.modbus_map.from_xml(name, pathlist)
        else:
            raise mbmap.ModbusMapError('No modbus map file provided during initialization')

    def close(self):

        pass

    def read(self, addr, count, op=None):

        if self.modbus_map is not None:
            return self.modbus_map.read(addr, count, op)
        else:
            raise ModbusClientError('No modbus map set for device')

    def write(self, addr, data):

        if self.modbus_map is not None:
            return self.modbus_map.write(addr, data)
        else:
            raise ModbusClientError('No modbus map set for device')

def __generate_crc16_table():
    ''' Generates a crc16 lookup table

    .. note:: This will only be generated once
    '''
    result = []
    for byte in range(256):
        crc = 0x0000
        for bit in range(8):
            if (byte ^ crc) & 0x0001:
                crc = (crc >> 1) ^ 0xa001
            else: crc >>= 1
            byte >>= 1
        result.append(crc)
    return result

__crc16_table = __generate_crc16_table()

def computeCRC(data):
    ''' Computes a crc16 on the passed in string. For modbus,
    this is only used on the binary serial protocols (in this
    case RTU).

    The difference between modbus's crc16 and a normal crc16
    is that modbus starts the crc value out at 0xffff.

    :param data: The data to create a crc16 of
    :returns: The calculated CRC
    '''
    crc = 0xffff
    for a in data:
        idx = __crc16_table[(crc ^ ord(a)) & 0xff];
        crc = ((crc >> 8) & 0xff) ^ idx
    swapped = ((crc << 8) & 0xff00) | ((crc >> 8) & 0x00ff)
    return swapped

def checkCRC(data, check):
    ''' Checks if the data matches the passed in CRC

    :param data: The data to create a crc16 of
    :param check: The CRC to validate
    :returns: True if matched, False otherwise
    '''
    return computeCRC(data) == check

