#!/usr/bin/env python3
# Software License Agreement (BSD)
#
# @author    Tony Baltovski <tbaltovski@clearpathrobotics.com>
# @copyright (c) 2020, Clearpath Robotics, Inc., All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that
# the following conditions are met:
# * Redistributions of source code must retain the above copyright notice, this list of conditions and the
#   following disclaimer.
# * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other materials provided with the distribution.
# * Neither the name of Clearpath Robotics nor the names of its contributors may be used to endorse or
#   promote products derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
# TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import can
import cantools
import os
import queue
import rospy
import threading
import time

from catkin.find_in_workspaces import find_in_workspaces
from diagnostic_msgs.msg import DiagnosticArray, DiagnosticStatus, KeyValue
from sensor_msgs.msg import BatteryState


class Battery:
    def __init__(self, id, db, db_data):
        self.id = id
        self.db = db
        self.db_data = db_data

    def check_for_level(self, key, level, message, level_to_check):
        if self.db_data[key] != 1:
            if level == level_to_check:
                message += (' ' + key)
            else:
                level = level_to_check
                message = key
        return level, message

    def process_msg(self, msg):
        s = self.db.decode_message(msg.arbitration_id, msg.data)

        for key in s:
            self.db_data.update({key[:-3]: s[key]})

    def get_bat_diag(self):
        bat_stat = DiagnosticStatus(name =('Vanguard Flex Battery ' + str(self.id)),
                                    level = DiagnosticStatus.OK, message = 'OK')

        for key in self.db_data:
            bat_stat.values.append(KeyValue(key, str(self.db_data[key])))

        # Warnings
        bat_stat.level, bat_stat.message = self.check_for_level('AFECommAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('BalanaceAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('BusVoltageMatchAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('ChargerNotMatchAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('ContactorAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('CurrentAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('EstopAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('LowVoltageAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('MaintenaceAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('OverVoltAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('StartVoltAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('TemperatureAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)
        bat_stat.level, bat_stat.message = self.check_for_level('ZeroCurrentAlarm', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.WARN)

        # Errors
        bat_stat.level, bat_stat.message = self.check_for_level('BMSAddrChangedProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('BMSMultipleAddressFault', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('BalanceProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('ChargerNotMatchProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('ContactorIssueProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('CurrentProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('EStopProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('FaultOouput', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('LowVoltageProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('MaintenProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('OverVoltageProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('PermanentFaultProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('StartVoltProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('TemperatureProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)
        bat_stat.level, bat_stat.message = self.check_for_level('ZeroCurrentProtect', bat_stat.level,
                                            bat_stat.message, DiagnosticStatus.ERROR)

        return bat_stat

    def get_soc(self):
        return self.db_data['SoC']

    def get_pack_voltage(self):
        return self.db_data['PackVoltagelvl']

    def get_pack_current(self):
        return self.db_data['PackCurrentlvl']

    def get_capacity(self):
        return self.db_data['DischargeCapacity']


def vanguard_bms_driver():
    node_name = 'vanguard_bms_driver'
    rospy.init_node(node_name)

    can_device = rospy.get_param('~can_device', 'can0')
    # Load the database.
    dbc_file = find_in_workspaces(project="vanguard_bms", path="config/Vanguard_Flex_Battery_4BAT_FormA.dbc")[0]
    db = cantools.database.load_file(dbc_file)

    # Get all the signals from the database and strip out the ID based ending.
    db_data = {}
    for i in range(len(db.messages)):
        for j in range(len(db.messages[i].signals)):
            db_data.update({db.messages[i].signals[j].name[:-3]: 0})

    q = queue.Queue()

    def can_recv():
        can_bus = can.interface.Bus(can_device, bustype='socketcan')
        while not rospy.is_shutdown():
            q.put(can_bus.recv())

    threading.Thread(target=can_recv, daemon=True).start()

    batteries_id = [3, 4, 5, 6]
    batteries = []
    for i in range(0, len(batteries_id)):
        batteries.append(Battery(batteries_id[i], db, db_data))

    pub_bat_msg = rospy.Publisher("battery_state", BatteryState, queue_size = 1)
    pub_diag = rospy.Publisher("/diagnostics", DiagnosticArray, queue_size = 1)

    last_time = time.time()
    while not rospy.is_shutdown():
        try:
            message = q.get()
            battery_node_id = int(message.arbitration_id) & int(15)
            batteries[battery_node_id - 3].process_msg(message)
        except KeyError:
            # print("Message ID not in database")
            pass
        if ((time.time() - last_time) >= 1):
            battery_msg = BatteryState()
            diag_array = DiagnosticArray()
            battery_msg.power_supply_technology = BatteryState.POWER_SUPPLY_TECHNOLOGY_LION
            battery_msg.present = True
            for i in range(0, len(batteries)):
                diag_array.status.append(batteries[i].get_bat_diag())
                battery_msg.current += batteries[i].get_pack_current()
                battery_msg.capacity += batteries[i].get_capacity()
                if i == 0:
                    battery_msg.percentage = batteries[0].get_soc()
                    battery_msg.voltage = batteries[0].get_pack_voltage()
            else:
                if batteries[i].get_soc() < battery_msg.percentage:
                    battery_msg.percentage = batteries[i].get_soc()
                    battery_msg.voltage = batteries[i].get_pack_voltage()
            battery_msg.header.stamp = rospy.Time.now()
            pub_bat_msg.publish(battery_msg)

            diag_array.header.stamp = battery_msg.header.stamp

            pub_diag.publish(diag_array)
            last_time = time.time()
    q.join()

if __name__ == '__main__':
    try:
        vanguard_bms_driver()
    except rospy.ROSInterruptException:
        pass
