#!/usr/bin/env python3

# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import asyncio
import psutil
import time
import os
import argparse
import logging
from datetime import datetime

from kasa import SmartStrip


class KasaPlugController:
    """Provides control of a device's charger.

    The device's charger must be plugged into one of the 3 outlets of a Kasa
    Smart
    Plug Power Strip (KP303). The outlet name must match the device's host name
    (this is intended to prevent inadvertently controlling the wrong device's
    charger).
    """

    def __init__(self, kasa_power_strip_ip: str):
        """
        Constructs a KasaPlugController to control the current device's
        charger.

        Args:
            kasa_power_strip_ip: IP of the Kasak Smart Plug Power Strip in which
                this device's charger is connected.
        """
        # The outlet name must match the device's host name.
        self.kasa_outlet_name = os.uname()[1].split('.')[0]

        # Create the event loop
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

        # Create the strip controller
        self.strip = SmartStrip(kasa_power_strip_ip)
        self.loop.run_until_complete(self.strip.update())

        self.closed = False

    def __del__(self):
        self.close()

    def turn_on(self):
        """Turns on this device's charger."""
        logging.info("Turning on the charger")
        for plug in self.strip.children:
            if plug.alias == self.kasa_outlet_name:
                self.loop.run_until_complete(plug.turn_on())
                return
        logging.error("Cannot find device to turn on")

    def turn_off(self):
        """Turns off this device's charger."""
        logging.info("Turning off the charger")
        for plug in self.strip.children:
            if plug.alias == self.kasa_outlet_name:
                self.loop.run_until_complete(plug.turn_off())

                battery = psutil.sensors_battery()
                while battery.power_plugged:
                    logging.info(
                        "Waiting for device to no longer be plugged in"
                    )
                    time.sleep(1)
                    battery = psutil.sensors_battery()
                return
        logging.error("Cannot find device to turn off")

    def discharge_to(self, level: int):
        """Discharges the battery until it reaches a target level.

        Args:
            level: The target battery level.
        """

        self.turn_off()

        battery = psutil.sensors_battery()

        while battery.percent > level:
            logging.info(
                f"Waiting to discharge to {level}%."
                f" Currently at {battery.percent}%"
            )

            # Perform arbitrary operations as fast as possible to burn
            # CPU and discharge faster.
            f_value = 0.81
            start = datetime.now()
            while (datetime.now() - start).total_seconds() < 10:
                f_value = f_value * 1.7272882
                f_value = f_value / 1.7272882

            battery = psutil.sensors_battery()

        logging.info(f"Discharge to {level}% complete")

    def charge_to(self, level: int):
        """Charges the battery until it reaches a target level.

        Args:
            level: The target battery level.
        """

        self.turn_on()

        battery = psutil.sensors_battery()

        while battery.percent < level:
            logging.info(
                f"Waiting to charge to {level}%."
                f" Currently at {battery.percent}%"
            )
            time.sleep(10)
            battery = psutil.sensors_battery()

        logging.info(f"Charge to {level}% complete")

    def charge_or_discharge_to(self, level: int):
        """Charges or discharges the battery until it reaches a target level.

        Leaves the charger in an unplugged state.

        Args:
            level: The target battery level.
        """
        battery = psutil.sensors_battery()
        if battery.percent < level:
            self.charge_to(level)
        elif battery.percent > level:
            self.discharge_to(level)
        else:
            logging.info(f"Battery is already at the target level {level}%")
        self.turn_off()

    def close(self):
        """Closes the message loop."""
        if self.closed:
            return
        self.closed = True
        self.loop.close()


def get_plug_controller(ip: str):
    return KasaPlugController(ip)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Controls kasa power switch connected to this device.'
    )
    parser.add_argument(
        "--kasa_power_strip_ip",
        required=True,
        help="IP address of the kasa power switch.",
    )
    parser.add_argument(
        "--charge_level", type=int, required=True, help="Desired charge level."
    )
    args = parser.parse_args()

    kasa_plug_controller = KasaPlugController(args.kasa_power_strip_ip)
    kasa_plug_controller.charge_to(args.charge_level)
    kasa_plug_controller.close()