import aiomqtt
import json
import asyncio
from datetime import datetime
from time import time
from pathlib import Path
from collections import namedtuple
from typing import Dict
from Cryptodome.Cipher import AES

BtleMeasurement = namedtuple("BtleMeasurement", ["time", "tracker", "address", "rssi", "tx_power"])
BtleDeviceMeasurement = namedtuple("BtleDeviceMeasurement", ["time", "device", "tracker", "rssi", "tx_power"])


class KnownRoomCsvLogger:
    """Logs known room measurements to be used later as training data for classifier"""

    def __init__(self, csv_file: Path):
        self.known_room = None

        if csv_file.exists():
            self.csv_file_handle = open(csv_file, "a")
        else:
            self.csv_file_handle = open(csv_file, "w")
            print(f"#time,device,tracker,rssi,tx_power,real_room", file=csv_file)

    def update_known_room(self, known_room: str):
        self.known_room = known_room

    def report_measure(self, m: BtleDeviceMeasurement):
        ignore_rooms = ("keins", "?", "none", "unknown")
        if self.known_room is None or self.known_room in ignore_rooms:
            return
        print(
            f"{m.time},{m.device},{m.tracker},{m.rssi},{m.tx_power},{self.known_room}",
            file=self.csv_file_handle,
        )


class DeviceDecoder:

    def __init__(self, irk_to_devicename: Dict[str, str], address_to_name: Dict[str, str]):
        "irk_to_devicename is dict with irk as a hex string, mapping to device name"
        self.irk_to_devicename = {bytes.fromhex(k): v for k, v in irk_to_devicename.items()}
        self.address_to_name = address_to_name

    def _resolve_rpa(rpa: bytes, irk: bytes) -> bool:
        """Compares the random address rpa to an irk (secret key) and return True if it matches"""
        assert len(rpa) == 6
        assert len(irk) == 16

        key = irk
        plain_text = b"\x00" * 16
        plain_text = bytearray(plain_text)
        plain_text[15] = rpa[3]
        plain_text[14] = rpa[4]
        plain_text[13] = rpa[5]
        plain_text = bytes(plain_text)

        cipher = AES.new(key, AES.MODE_ECB)
        cipher_text = cipher.encrypt(plain_text)
        return cipher_text[15] == rpa[0] and cipher_text[14] == rpa[1] and cipher_text[13] == rpa[2]

    def _addr_to_bytes(addr: str) -> bytes:
        """Converts a bluetooth mac address string with semicolons to bytes"""
        str_without_colons = addr.replace(":", "")
        bytearr = bytearray.fromhex(str_without_colons)
        bytearr.reverse()
        return bytes(bytearr)

    def addr_to_bytes(addr: str) -> bytes:
        """Converts a bluetooth mac address string with semicolons to bytes"""
        str_without_colons = addr.replace(":", "")
        bytearr = bytearray.fromhex(str_without_colons)
        bytearr.reverse()
        return bytes(bytearr)

    def decode(self, addr: str):
        """addr is a bluetooth address as a string  e.g. 4d:24:12:12:34:10"""
        for irk, name in self.irk_to_devicename.items():
            if self.resolve_rpa(self.addr_to_bytes(addr), irk):
                return name
        return self.addr_to_name.get(addr, None)


server = "homeassistant.fritz.box"
username = "my_btmonitor"
password = "8aBIAC14jaKKbla"


async def collect_data_from_mqtt_into_csv():
    now = datetime.now()
    with open(f"logfile_{now}.csv", "w") as csv_file:
        print(f"# time,device,room,rssi,tx_power,real_room", file=csv_file)
        async with aiomqtt.Client(hostname=server, username=username, password=password) as client:
            real_room = "?"
            await client.subscribe("my_btmonitor/#")
            async for message in client.messages:
                current_time = time()
                topic = message.topic
                if topic.value == "my_btmonitor/real_room":
                    print(f"Changing real room from {real_room} to {message.payload}")
                    real_room = message.payload.decode()
                else:
                    splitted_topic = message.topic.value.split("/")
                    if splitted_topic[0] == "my_btmonitor" and splitted_topic[1] == "devices":
                        device = splitted_topic[2]
                        room = splitted_topic[3]
                        msg_json = json.loads(message.payload)
                        rssi = msg_json.get("rssi", -1)
                        tx_power = msg_json.get("tx_power", -1)
                        if real_room is not None and real_room != "keins":
                            print(
                                f"{current_time},{device},{room},{rssi},{tx_power},{real_room}",
                                file=csv_file,
                            )
                            print(f"{current_time},{device},{room},{rssi},{tx_power},{real_room}")


if __name__ == "__main__":
    asyncio.run(collect_data_from_mqtt_into_csv())