import argparse
import can
import cantools
import json
import paho.mqtt.client as mqtt
import sys
import time
from pathlib import Path
from typing import Iterable

def prepare_j1939_db(db):
    # Create a J1939 PDU1 and PDU2 database (with matching PGN ID masks)
    db_pdu1 = cantools.database.Database(frame_id_mask=0x3FF0000)
    db_pdu2 = cantools.database.Database(frame_id_mask=0x3FFFF00)

    for db_msg in db.messages:
        if cantools.j1939.is_pdu_format_1(cantools.j1939.frame_id_unpack(db_msg.frame_id).pdu_format):
            db_pdu1.messages.append(db_msg)
        else:
            db_pdu2.messages.append(db_msg)

    db_pdu1.refresh()
    db_pdu2.refresh()

    return db_pdu1, db_pdu2

def can2mqtt(bus: Iterable[can.Message], source_name: str, db, mqtt_host: str, mqtt_port: int, use_system_time=False):

    # Connect to MQTT broker
    print(f"Connecting to MQTT broker: {mqtt_host}:{mqtt_port}")
    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
    client.connect(host=mqtt_host, port=mqtt_port)
    client.loop_start()

    # If DBC ProtocolType is J1939 prepare for PGN matching
    protocol_type = db.dbc.attributes.get("ProtocolType", None)
    if protocol_type and protocol_type.value == "J1939":
        db_pdu1, db_pdu2 = prepare_j1939_db(db)

    # Forward CAN-bus to MQTT
    topics = []
    print("\nPress \"ctrl + c\" to quit\n\nMQTT topics:")
    try:
        # Receive CAN-bus
        for msg in bus:

            # If DBC ProtocolType is J1939, match based on J1939 PGN instead of 29 bit CAN ID
            if protocol_type == "J1939":
                if cantools.j1939.is_pdu_format_1(cantools.j1939.frame_id_unpack(msg.arbitration_id).pdu_format):
                    db = db_pdu1
                else:
                    db = db_pdu2

            # Get message name from DB
            try:
                db_msg = db.get_message_by_frame_id(msg.arbitration_id)
            except KeyError:
                continue

            # Decoder CAN-bus frame
            signals = db.decode_message(db_msg.name, msg.data, decode_choices=False)

            # Add timestamp as integer in ms
            if use_system_time is True:
                signals["timestamp"] = int(time.time()*1000)
            else:
                signals["timestamp"] = int(msg.timestamp * 1000)

            # Construct topic
            topic = f"{source_name}/{db_msg.name}"
            if topic not in topics:
                topics.append(topic)
                print(f"- {topic}")

            # MQTT publish
            client.publish(topic, json.dumps(signals))
    except KeyboardInterrupt:
        pass

    client.disconnect()


def main():

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--itf", type=str, default="csscan_serial", help="CAN-bus interface, e.g. csscan_serial:COM1")
    parser.add_argument("--file", type=str, default=None, help="CAN-bus log file, e.g. 00000001.MF4")
    parser.add_argument("--dbc", type=Path, default="db.dbc", help="CAN-bus database file path")
    parser.add_argument("--host", type=str, default="localhost", help="MQTT broker host")
    parser.add_argument("--port", type=int, default=1883, help="MQTT broker port")
    args = parser.parse_args()

    # Load CAN-bus database
    if not args.dbc.is_file():
        print(f"DBC \"{args.dbc}\" not found")
        sys.exit(1)

    print(f"Loading CAN-bus database: \"{args.dbc}\"")
    db = cantools.database.load_file(args.dbc)

    if not args.itf and not args.file:
        print("Provide interface or file")
        sys.exit(1)

    elif args.file:
        # CAN-bus from file
        log_file = Path(args.file)

        if not log_file.is_file():
            print(f"Cannot open log file \"{log_file}\"")
            sys.exit(1)

        # Simulate bus using MessageSync
        if log_file.suffix.lower() in (".mf4", ".mfe", ".mfc", ".mfm"):
            reader = LogFileReader(log_file)
        else:
            reader = can.LogReader(log_file)

        reader = can.MessageSync(reader)
        can2mqtt(bus=reader, source_name="file", db=db, mqtt_host=args.host, mqtt_port=args.port, use_system_time=True)

    elif args.itf:
        # CAN-bus from interface

        interface_channel = args.itf.split(":")
        itf_name = interface_channel[0]
        itf_chn = None if len(interface_channel) < 2 else interface_channel[1]

        # Find all CAN-bus interface channels
        itf_chns = [x["channel"] for x in can.detect_available_configs(itf_name)]

        # If no interface channel provided, default to first found
        if itf_chn is None and len(itf_chns) > 0:
            itf_chn = itf_chns[0]

        # Check if provided interface channel is available
        if itf_chn not in itf_chns:
            print(f"CAN-bus channel \"{itf_chn}\" not available")
            if len(itf_chns) > 0:
                print(f"Available channels: {', '.join(itf_chns)}")
            sys.exit(-1)

        # Open CAN-bus interface
        print(f"Starting CAN-bus interface: {itf_name}:{itf_chn}")
        with can.Bus(interface=itf_name, channel=itf_chn) as bus:
            can2mqtt(bus=bus, source_name=itf_chn, db=db, mqtt_host=args.host, mqtt_port=args.port)


if __name__ == '__main__':
    main()
