#! /usr/bin/env python
"""
This tool is intended to repair XBee modules that have a miscalibrated
frequency alignment. See the included Readme for detailed usage information.
"""
from __future__ import annotations

import argparse
import binascii
import ctypes
import dataclasses
import logging
import re
import sys
import time
import zipfile
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union, cast

import tomli
from digi.xbee.devices import RemoteDigiMeshDevice, XBeeDevice
from digi.xbee.models.address import XBee64BitAddress
from digi.xbee.exception import ATCommandException
from digi.xbee.serial import FlowControl

log = logging.getLogger("xb_repair")
log.setLevel(logging.DEBUG)
digi_log = logging.getLogger("digi")
digi_log.setLevel(logging.DEBUG)

here = Path(__file__).resolve().parent

VERSION = (1, 0)

FirmwareTuple = Tuple[int, Path]  # Used to represent a firmware version in some places


def main():
    parser = argparse.ArgumentParser(
        description=cast(str, __doc__) + f"\n\nv{VERSION[0]}.{VERSION[1]}",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    subparsers = parser.add_subparsers(dest="command")
    parser.add_argument(
        "--configuration", "-C", type=Path, help="Path to a configuration zip file."
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true", help="enable verbose logging"
    )
    parser.set_defaults(fn=None)

    def add_port_args(p: argparse.ArgumentParser, fastplease: bool):
        p.add_argument(
            "--port",
            "-p",
            required=True,
            help="The name of the port to use to communicate with the XBee module.",
        )
        p.add_argument(
            "--baud",
            "-b",
            type=int,
            default=115200,
            help=(
                "The baud rate the XBee module is configured to."
                + (
                    " It is recommended to use a baud of 115200 (set ATBD to 7) or higher."
                    if fastplease
                    else ""
                )
            ),
        )

    search = subparsers.add_parser("search", help=main_search.__doc__)
    add_port_args(search, True)
    search.add_argument(
        "--output",
        "-o",
        type=Path,
        help="The name of a file to save the list of discovered radios.",
        required=True,
    )
    search.add_argument(
        "--count",
        "-c",
        default=None,
        type=int,
        help="The number of radios expected to be found. Discovery will stop early if this number is reached.",
    )
    search.set_defaults(fn=main_search)

    repair = subparsers.add_parser("repair", help=main_repair.__doc__)
    add_port_args(repair, True)
    repair.add_argument(
        "--input",
        "-i",
        required=True,
        type=Path,
        help="A file containing a list of radios as generated by the search step.",
    )
    repair.set_defaults(fn=main_repair)

    repair_local = subparsers.add_parser("repair-local", help=main_repair_local.__doc__)
    add_port_args(repair_local, False)
    repair_local.set_defaults(fn=main_repair_local)

    args = parser.parse_args()

    if args.command is None:
        parser.print_help()
        sys.exit(1)

    log_file = here / f"xb_{args.command}.log"
    configure_logging(logging.StreamHandler(sys.stderr), args.verbose)
    configure_logging(logging.FileHandler(log_file, "w"), True)

    log.debug("args: %s", args)

    try:
        args.configuration = Config.from_file(args.configuration)
        args.fn(args)
    except:
        log.exception("Fatal error")
        log.error(
            "An error has occurred. See the detailed log at %s for more details.",
            log_file,
        )


def main_search(args: argparse.Namespace):
    """Search over-the air for modules to repair.

    This step uses a serially connected module to search for affected
    modules nearby. The modules must be powered up, and the local module
    needs to be configured with the same network settings as the modules
    to fix. A network discovery will be done at several different
    frequency offsets, and all radios found will be saved in the output
    file.
    """
    config: Config = args.configuration
    host = open_radio(args.port, args.baud, config, True)

    net = host.get_network()
    net.set_discovery_options({})
    net.set_discovery_timeout(config.discovery_timeout)

    with open(args.output, "w") as output:
        discovered: Set[DiscoveredModule] = set()
        skipped: Set[bytes] = set()

        for offset in config.search_values:
            log.info("Searching at frequency offset %d", offset)
            set_f_percent(host, offset)

            net.start_discovery_process()
            time.sleep(1)
            assert (
                net.is_discovery_running()
            ), "Error: Couldn't start network discovery. This can happen if the script exits early. Try waiting ~15 seconds or resetting the XBee."
            while net.is_discovery_running():
                pass

            devices = net.get_devices()
            log.info("received %d responses", len(devices))

            for device in devices:
                try:
                    device.set_sync_ops_timeout(config.remote_timeout)
                    addr64 = device.get_64bit_addr()
                    log.info("  address: %s", addr64)

                    already_discovered = addr64.address in (
                        d.address for d in discovered
                    )
                    already_skipped = addr64.address in (address for address in skipped)
                    if already_discovered or already_skipped:
                        log.info(
                            "    This radio was found in a previous scan, skipping..."
                        )
                        continue

                    hv = read_hv(device)
                    log.info("    HV: 0x%02X", hv)
                    reg = int.from_bytes(device.get_parameter("R?"), "big")
                    log.info("    R?: 0x%02X", reg)

                    try:
                        config.for_hw(hv, reg)
                    except KeyError:
                        log.info("    this radio is not supported, skipping...")
                        skipped.add(bytes(addr64.address))
                        continue

                    freq_offset = from_i16(device.get_parameter("F%"))
                    log.info("    frequency offset: %d", freq_offset)
                    discovered.add(
                        DiscoveredModule(
                            address=bytes(addr64.address),
                            freq_offset=freq_offset,
                            hv=hv,
                            reg=reg,
                        )
                    )

                except:
                    log.error("Failed to retrieve information on %s", device)
                    log.debug("Failed to retreive information on %s", exc_info=True)

            if args.count is not None and len(discovered) >= args.count:
                log.info("Found all %d expected modules", args.count)
                break
        else:
            log.info("Found %d modules")

        for d in discovered:
            output.write(f"{d.to_str()}\n")


def main_repair(args: argparse.Namespace):
    """Repair radios from a previous over-the-air search.

    This step takes the list of radios from the input file and repairs
    them one by one. Radios will be skipped if they do not need to be
    repaired.
    """
    config: Config = args.configuration
    host = open_radio(args.port, args.baud, config, True)
    targets = DiscoveredModule.from_file(args.input)

    for target in targets:
        try:
            repair_remote_radio(host, target, config)
        except:
            log.error("Update failed on %s", target)
            log.debug("Update failed on %s", target, exc_info=True)


def main_repair_local(args: argparse.Namespace):
    """Repair a radio over the serial port.

    This step. will check the frequency offset on the specified radio,
    and repair it if needed.
    """
    config: Config = args.configuration
    radio = open_radio(args.port, args.baud, config, False)
    repair_radio(radio, config)


def configure_logging(handler: logging.Handler, verbose: bool):
    """Set up a handler for logging."""
    handler.setLevel(logging.DEBUG if verbose else logging.INFO)
    formatter = logging.Formatter("%(asctime)s:%(name)20s:%(levelname)5s: %(message)s")
    handler.setFormatter(formatter)
    digi_log.addHandler(handler)
    log.addHandler(handler)


def install_fw(radio: XBeeDevice, firmware: FirmwareTuple):
    """Install firmware on a radio"""
    _, fw = firmware
    fw = str(here / fw)
    log.debug("Load %s on %s", fw, radio)
    radio.update_firmware(str(fw))


def repair_radio(radio: XBeeDevice, config: Config, host: Optional[XBeeDevice] = None):
    """Go  through the repair process on a radio.

    If repairing a remote radio, 'host' should be the local XBeeDevice.
    """
    version = int.from_bytes(radio.get_parameter("VR"), "big")
    hv = read_hv(radio)
    reg = int.from_bytes(radio.get_parameter("R?"), "big")
    log.info("Starting version: %04X, HV: %X, R?: %X", version, hv, reg)

    hw_config = config.for_hw(hv, reg)
    fp_initial = from_i16(radio.get_parameter("F%"))
    log.info("Starting frequency offset: %d", fp_initial)

    # First, fix frequency if needed
    if abs(fp_initial) > hw_config.threshold:
        if version != hw_config.repair_fw[0]:
            log.info("Load repair firmware on %s...", radio)
            install_fw(radio, hw_config.repair_fw)

        log.info("Set Frequency offset to %d on target", hw_config.target)
        set_f_percent(radio, hw_config.target)

        if host is not None:
            log.info("Move host back to frequency offset %d", hw_config.target)
            set_f_percent(host, hw_config.target)
    else:
        log.info(
            "Frequency offset on %s is %d: This radio doesn't need to be repaired",
            radio.get_64bit_addr(),
            fp_initial,
        )

    # Whether we fixed the frequency or not, don't leave repair firmware
    # on the radio
    version = int.from_bytes(radio.get_parameter("VR"), "big")
    if version == hw_config.repair_fw[0]:
        log.info("Load release firmware on %s...", radio)
        install_fw(radio, hw_config.release_fw)

    fp_ending = from_i16(radio.get_parameter("F%"))
    log.info("Frequency offset changed from 0x%X to 0x%X", fp_initial, fp_ending)


def repair_remote_radio(host: XBeeDevice, target: DiscoveredModule, config: Config):
    log.debug("Attempting repair on %r", target)
    address = XBee64BitAddress(bytearray(target.address))
    remote = RemoteDigiMeshDevice(host, address)
    remote.set_sync_ops_timeout(config.remote_timeout)

    # Check the recorded frequency and 0, just in case the target
    # has already been fixed.
    freq = 0
    for freq in (target.freq_offset, 0):
        log.info("Look for target %s at frequency offset %d", target, freq)
        set_f_percent(host, freq)
        try:
            fp_initial = from_i16(remote.get_parameter("F%"))
            break
        except ATCommandException:
            log.info(
                "Couldn't communicate with target %s at frequency offset %d",
                target,
                freq,
            )
    else:
        log.error(
            "Couldn't find target radio %s at frequency %d or 0",
            target,
            target.freq_offset,
        )
        return

    log.info(
        "Found target %s on frequency %s, target has frequency offset %d",
        target,
        freq,
        fp_initial,
    )
    # Get on the frequency and do the repair
    if fp_initial != freq:
        log.info("Move host radio to offset %d to match target", fp_initial)
        set_f_percent(host, fp_initial)

    repair_radio(remote, config, host=host)


def set_f_percent(radio: XBeeDevice, offset: int):
    """Set F% on a radio and apply it immediately."""
    f_percent_value = ctypes.c_uint16(offset).value
    radio.set_parameter("F%", f_percent_value.to_bytes(2, "big"))
    radio.set_parameter("WR", b"")
    radio.set_parameter("WP", b"")
    radio.set_parameter("FR", b"")
    time.sleep(3)


def open_radio(
    port: str,
    baud: int,
    config: Config,
    load_fw: bool,
):
    """Open a local radio, and optionally make sure it has repair fw"""
    log.info("Connecting to XBee on %s", port)
    radio = XBeeDevice(port, baud, flow_control=FlowControl.HARDWARE_RTS_CTS)
    radio.set_sync_ops_timeout(config.host_timeout)
    radio.open(force_settings=True)

    version = int.from_bytes(radio.get_parameter("VR"), "big")
    hv = read_hv(radio)
    reg = int.from_bytes(radio.get_parameter("R?"), "big")

    host_firmware = config.for_hw(hv, reg).host_fw
    if load_fw and version != host_firmware[0]:
        log.info(
            "Loading repair firmware %X on the radio at %s so that it can act as a host.",
            host_firmware[0],
            port,
        )
        install_fw(radio, host_firmware)

    return radio


def from_i16(val: Union[int, bytes]) -> int:
    """Turn bytes from an AT command response into a uint16."""
    if not isinstance(val, int):
        val = int.from_bytes(val, "big")
    return ctypes.c_int16(val).value


def read_hv(radio: XBeeDevice):
    """Read HV on a radio"""
    hv = radio.get_hardware_version()
    if hv is not None:
        return hv.code
    else:
        return int.from_bytes(radio.get_parameter("HV"), "big") >> 8


@dataclasses.dataclass(frozen=True)
class DiscoveredModule:
    """Represents a module discovered in the search step."""

    address: bytes
    freq_offset: int
    hv: int
    reg: int

    def __str__(self) -> str:
        return binascii.hexlify(self.address).decode()

    def to_str(self):
        """String representation for writing to a file."""
        addr = binascii.hexlify(self.address).decode()
        return f"{addr}:{self.freq_offset:04X}:{self.hv:02X}:{self.reg:02X}"

    @classmethod
    def from_str(cls, s: str) -> DiscoveredModule:
        """Parse a DiscoveredModule from its to_str representation."""
        match = re.match(
            r"^(?P<addr>[0-9a-fA-F]{16}):(?P<freq>-?[0-9A-F]{1,5}):(?P<hv>[0-9A-F]{2}):(?P<reg>[0-9A-F]{2})$",
            s,
        )
        assert match is not None, f"Invalid radio entry: {s}"
        address = bytes.fromhex(match.group("addr"))
        freq_offset = int(match.group("freq"), 16)
        hv = int(match.group("hv"), 16)
        reg = int(match.group("reg"), 16)

        return DiscoveredModule(
            address=address,
            freq_offset=freq_offset,
            hv=hv,
            reg=reg,
        )

    @classmethod
    def from_file(cls, path: Path) -> List[DiscoveredModule]:
        """Read a list of DiscoveredModules, one per line, from a file."""
        return [cls.from_str(line.strip()) for line in path.read_text().splitlines()]


@dataclasses.dataclass
class Config:
    """Configuration from config.zip"""

    search_values: List[int]

    host_timeout: int
    remote_timeout: int
    discovery_timeout: int

    hv_reg_firmware: Dict[str, Dict[str, HardwareConfig]]
    _tempdir: Optional[TemporaryDirectory[str]] = None

    @classmethod
    def from_file(cls, file: Optional[Path]) -> Config:
        """Load either a config toml or zip file."""
        if file is None:
            for default in [here / "config.toml", here / "xbee_repair_config.zip"]:
                if default.exists():
                    log.debug("Using default config file %s", default)
                    return Config.from_file(default)
                else:
                    log.debug("Check for default config file %d: not found", default)
            raise RuntimeError("No config specified and no default found.")
        else:
            if file.suffix == ".zip":
                return Config.from_zip(file)
            else:
                return Config.from_config(file)

    @classmethod
    def from_zip(cls, file: Path) -> Config:
        """Load a config zip file"""
        tempdir = TemporaryDirectory(prefix="xbrep")
        log.debug("Extract %s to %s", file, tempdir)
        with zipfile.ZipFile(file) as z:
            z.extractall(tempdir.name)
        config = Config.from_config(Path(tempdir.name) / "config.toml")
        config._tempdir = tempdir
        return config

    @classmethod
    def from_config(cls, file: Path) -> Config:
        """Load a config toml file"""
        log.debug("Load config from %s", file)
        config_dir = Path(file).resolve().parent
        with open(file, "rb") as f:
            data: Any = tomli.load(f)  # type: ignore

        hv_reg_fw: Dict[str, Dict[str, Dict[str, List[Any]]]] = data["hv_reg_firmware"]
        for hv, reg_fw in hv_reg_fw.items():
            for reg, fw in reg_fw.items():
                fw["host_fw"][1] = config_dir / fw["host_fw"][1]
                fw["repair_fw"][1] = config_dir / fw["repair_fw"][1]
                fw["release_fw"][1] = config_dir / fw["release_fw"][1]
                hv_reg_fw[hv][reg] = HardwareConfig(**fw)  # type: ignore

        config = Config(**data)
        log.debug("Config: %s", config)

        return config

    def for_hw(self, hv: int, reg: int) -> HardwareConfig:
        """Get the hardware config for a given hardware."""
        hv_key = format(hv, "02X")
        reg_key = format(reg, "02X")
        try:
            return self.hv_reg_firmware[hv_key][reg_key]
        except KeyError:
            return self.hv_reg_firmware[hv_key]["*"]

    def hw_configs(self) -> Iterable[HardwareConfig]:
        for reg_fw in self.hv_reg_firmware.values():
            for fw in reg_fw.values():
                yield fw


@dataclasses.dataclass
class HardwareConfig:
    desc: str
    threshold: int
    target: int
    host_fw: FirmwareTuple
    repair_fw: FirmwareTuple
    release_fw: FirmwareTuple


if __name__ == "__main__":
    main()
