#!/usr/bin/env python3
import argparse
import asyncio
import base64
import sys
import time
from contextlib import aclosing
from typing import List

import broadlink
from broadlink.const import DEFAULT_PORT
from broadlink.remote import CapturedSignal, data_to_pulses, pulses_to_data

TIMEOUT = 30


def auto_int(x):
    return int(x, 0)


def format_pulses(pulses: List[int]) -> str:
    """Format pulses."""
    return " ".join(
        f"+{pulse}" if i % 2 == 0 else f"-{pulse}"
        for i, pulse in enumerate(pulses)
    )


def parse_pulses(data: List[str]) -> List[int]:
    """Parse pulses."""
    return [abs(int(s)) for s in data]


def show(signal: CapturedSignal) -> None:
    """Print a captured signal in every format and save it if asked."""
    raw_fmt = signal.packet.hex()
    base64_fmt = base64.b64encode(signal.packet).decode('ascii')
    pulse_fmt = format_pulses(signal.pulses)

    print("Packet found!")
    if signal.frequency_mhz:
        print("Frequency: {}MHz".format(signal.frequency_mhz))
    print("Raw:", raw_fmt)
    print("Base64:", base64_fmt)
    print("Pulses:", pulse_fmt)

    if args.learnfile:
        print("Saving to {}".format(args.learnfile))
        with open(args.learnfile, "w") as text_file:
            text_file.write(pulse_fmt if args.durations else raw_fmt)


async def listen(window) -> int:
    """Drain a capture window, printing each signal; return how many."""
    heard = 0
    async with aclosing(window) as signals:
        async for signal in signals:
            heard += 1
            show(signal)
    return heard


parser = argparse.ArgumentParser(fromfile_prefix_chars='@')
parser.add_argument("--device", help="device definition as 'type host mac'")
parser.add_argument("--type", type=auto_int, default=0x2712, help="type of device")
parser.add_argument("--host", help="host address")
parser.add_argument("--mac", help="mac address (hex reverse), as used by python-broadlink library")
parser.add_argument("--temperature", action="store_true", help="request temperature from device")
parser.add_argument("--humidity", action="store_true", help="request humidity from device")
parser.add_argument("--energy", action="store_true", help="request energy consumption from device")
parser.add_argument("--check", action="store_true", help="check current power state")
parser.add_argument("--checknl", action="store_true", help="check current nightlight state")
parser.add_argument("--turnon", action="store_true", help="turn on device")
parser.add_argument("--turnoff", action="store_true", help="turn off device")
parser.add_argument("--turnnlon", action="store_true", help="turn on nightlight on the device")
parser.add_argument("--turnnloff", action="store_true", help="turn off nightlight on the device")
parser.add_argument("--switch", action="store_true", help="switch state from on to off and off to on")
parser.add_argument("--send", action="store_true", help="send command")
parser.add_argument("--sensors", action="store_true", help="check all sensors")
parser.add_argument("--learn", action="store_true", help="learn command")
parser.add_argument("--rflearn", action="store_true", help="rf scan learning")
parser.add_argument("--frequency", type=float, help="specify radiofrequency for learning")
parser.add_argument("--learnfile", help="save learned command to a specified file")
parser.add_argument("--window", type=float, default=TIMEOUT,
                    help="seconds to keep listening while learning (default %(default)s)")
parser.add_argument("--keep", action="store_true",
                    help="keep listening for the whole window and print every code heard")
parser.add_argument("--repeat", type=int, default=0,
                    help="with --send --durations: extra transmissions after the first")
parser.add_argument("--durations", action="store_true",
                    help="use durations in micro seconds instead of the Broadlink format")
parser.add_argument("--convert", action="store_true", help="convert input data to durations")
parser.add_argument("--joinwifi", nargs=2, help="Args are SSID PASSPHRASE to configure Broadlink device with")
parser.add_argument("data", nargs='*', help="Data to send or convert")
args = parser.parse_args()


async def main():
    dev = None

    if args.device:
        values = args.device.split()
        devtype = int(values[0], 0)
        host = values[1]
        mac = bytearray.fromhex(values[2])
    elif args.mac:
        devtype = args.type
        host = args.host
        mac = bytearray.fromhex(args.mac)

    if args.host or args.device:
        dev = broadlink.gendevice(devtype, (host, DEFAULT_PORT), mac)
        await dev.auth()

    if args.joinwifi:
        await broadlink.setup(args.joinwifi[0], args.joinwifi[1], 4)

    if args.convert:
        data = bytearray.fromhex(''.join(args.data))
        pulses = data_to_pulses(data)
        print(format_pulses(pulses))
    if args.temperature:
        print(await dev.check_temperature())
    if args.humidity:
        print(await dev.check_humidity())
    if args.energy:
        print(await dev.get_energy())
    if args.sensors:
        data = await dev.check_sensors()
        for key in data:
            print("{} {}".format(key, data[key]))
    if args.send:
        data = (
            pulses_to_data(parse_pulses(args.data), repeat=args.repeat)
            if args.durations
            else bytes.fromhex(''.join(args.data))
        )
        await dev.send_data(data)
    if args.learn or (args.learnfile and not args.rflearn):
        print("Learning...")
        heard = await listen(dev.capture(window=args.window, stop_after_first=not args.keep))
        if not heard:
            print("No data received...")
            sys.exit(1)
    if args.check:
        if await dev.check_power():
            print('* ON *')
        else:
            print('* OFF *')
    if args.checknl:
        if await dev.check_nightlight():
            print('* ON *')
        else:
            print('* OFF *')
    if args.turnon:
        await dev.set_power(True)
        if await dev.check_power():
            print('== Turned * ON * ==')
        else:
            print('!! Still OFF !!')
    if args.turnoff:
        await dev.set_power(False)
        if await dev.check_power():
            print('!! Still ON !!')
        else:
            print('== Turned * OFF * ==')
    if args.turnnlon:
        await dev.set_nightlight(True)
        if await dev.check_nightlight():
            print('== Turned * ON * ==')
        else:
            print('!! Still OFF !!')
    if args.turnnloff:
        await dev.set_nightlight(False)
        if await dev.check_nightlight():
            print('!! Still ON !!')
        else:
            print('== Turned * OFF * ==')
    if args.switch:
        if await dev.check_power():
            await dev.set_power(False)
            print('* Switch to OFF *')
        else:
            await dev.set_power(True)
            print('* Switch to ON *')
    if args.rflearn:
        if args.frequency:
            frequency = args.frequency
            print("Press the button you want to learn, a short press...")
        else:
            await dev.sweep_frequency()
            print("Detecting radiofrequency, press and hold the button to learn...")

            start = time.time()
            while time.time() - start < TIMEOUT:
                await asyncio.sleep(1)
                locked, frequency = await dev.check_frequency()
                if locked:
                    break
            else:
                print("Radiofrequency not found")
                await dev.cancel_sweep_frequency()
                sys.exit(1)

            print("Radiofrequency detected: {}MHz".format(frequency))
            print("You can now let go of the button")

            input("Press enter to continue...")

            print("Press the button again, now a short press.")

        heard = await listen(
            dev.capture_rf(window=args.window, frequency=frequency, stop_after_first=not args.keep)
        )
        if not heard:
            print("No data received...")
            sys.exit(1)


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