import glob
import threading
import time
import logging
import os
import re

import gi

from powersupply.platform_internals import AXP803, PMI8998

gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, GLib, GObject, Gio, GdkPixbuf

gi.require_version('Handy', '1')
from gi.repository import Handy

logging.basicConfig(level=logging.DEBUG)


class ExternalBattery:
    def __init__(self):
        self.name = None
        self.present = False
        self.status = None
        self.capacity = None
        self.voltage = None
        self.current = None
        self.health = None
        self.temp = None


class PowerStatus:
    def __init__(self):
        self.battery_present = None
        self.battery_charging = None
        self.battery_capacity = None
        self.battery_voltage = None
        self.battery_current = None
        self.battery_health = None
        self.battery_temp = None

        self.supply_type = None

        self.usb_present = None
        self.usb_current_limit = None
        self.usb_current = None
        self.usb_voltage = None
        self.usb_type = None

        self.ac_present = None

        self.typec = None
        self.typec_revision = None
        self.typec_pd_revision = None
        self.typec_pr = None
        self.typec_dr = None
        self.typec_operation_mode = None

        self.ext = []


class DataPoller(threading.Thread):
    def __init__(self, callback):
        threading.Thread.__init__(self)
        self.callback = callback

        self.regex_enum = re.compile(r'\[([^\]]+)\]')

        self.opmode = {
            'default': 'Default',
            '1.5A': '1.5A no PD',
            '3A': '3A no PD',
            '3.0A': '3.0A no PD',
            'usb_power_delivery': 'PD',
        }

    def run(self):
        supplies = self.get_supply_types()
        print("Supplies found: {}".format(", ".join(supplies.keys())))

        battery_present = os.path.join(supplies['bat'][0], 'present')
        battery_status = os.path.join(supplies['bat'][0], 'status')
        battery_health = os.path.join(supplies['bat'][0], 'health')
        battery_capacity = os.path.join(supplies['bat'][0], 'capacity')
        battery_voltage = os.path.join(supplies['bat'][0], 'voltage_now')
        battery_temp = os.path.join(supplies['bat'][0], 'temp')
        battery_current = os.path.join(supplies['bat'][0], 'current_now')

        if not os.path.isfile(battery_present):
            battery_present = None
        if not os.path.isfile(battery_status):
            battery_status = None
        if not os.path.isfile(battery_health):
            battery_health = None
        if not os.path.isfile(battery_capacity):
            battery_capacity = None
        if not os.path.isfile(battery_voltage):
            battery_voltage = None
        if not os.path.isfile(battery_temp):
            battery_temp = None
        if not os.path.isfile(battery_current):
            battery_current = None

        supply_type = 'USB'

        if 'usb' in supplies:
            usb_present = os.path.join(supplies['usb'][0], 'present')
            if not os.path.isfile(usb_present):
                usb_present = os.path.join(supplies['usb'][0], 'online')
            if not os.path.isfile(usb_present):
                usb_present = None

            usb_current_limit = os.path.join(supplies['usb'][0], 'input_current_limit')
            if not os.path.isfile(usb_current_limit):
                usb_current_limit = os.path.join(supplies['usb'][0], 'current_max')
            if not os.path.isfile(usb_current_limit):
                usb_current_limit = None

            usb_voltage = os.path.join(supplies['usb'][0], 'voltage_now')
            if not os.path.isfile(usb_voltage):
                usb_voltage = None

            usb_current = os.path.join(supplies['usb'][0], 'current_now')
            if not os.path.isfile(usb_current):
                usb_current = None

            usb_status = os.path.join(supplies['usb'][0], 'status')
            if not os.path.isfile(usb_status):
                usb_status = None

            usb_type = os.path.join(supplies['usb'][0], 'usb_type')
            if not os.path.isfile(usb_type):
                usb_type = None

        elif 'ac' in supplies:
            supply_type = 'AC'

        if 'ac' in supplies:
            ac_present = os.path.join(supplies['ac'][0], 'online')

        platform = None
        if os.path.isfile('/proc/device-tree/compatible'):
            compatible = self.read_sysfs_str('/proc/device-tree/compatible')
            if 'sun50i-a64' in compatible:
                # Probably Allwinner A64 with AXP803
                platform = AXP803()
        if os.path.isdir('/sys/class/power_supply/pmi8998_charger'):
            platform = PMI8998()

        while True:
            ps = PowerStatus()
            ps.supply_type = supply_type
            ps.battery_present = self.read_sysfs(battery_present) == 1 if battery_present is not None else True
            ps.usb_present = 'usb' in supplies and self.read_sysfs(usb_present) == 1
            ps.ac_present = 'ac' in supplies and self.read_sysfs(ac_present) == 1
            ps.battery_health = self.read_sysfs_str(battery_health)
            ps.battery_capacity = self.read_sysfs(battery_capacity)
            bat_volt = self.read_sysfs(battery_voltage) / 1000000.0
            ps.battery_voltage = bat_volt if bat_volt > 0 else None
            ps.battery_current = self.read_sysfs(battery_current) / 1000000.0

            if 'usb' in supplies:
                ucl = self.read_sysfs(usb_current_limit)
                if ucl:
                    ps.usb_current_limit = ucl / 1000000.0
                uv = self.read_sysfs(usb_voltage)
                if uv:
                    ps.usb_voltage = uv / 1000000.0
                uc = self.read_sysfs(usb_current)
                if uc:
                    ps.usb_current = uc / 1000000.0
                ut = self.read_sysfs_enum(usb_type)
                if ut:
                    ps.usb_type = ut

            if battery_temp:
                bt = self.read_sysfs(battery_temp)
                if bt is None:
                    ps.battery_temp = None
                else:
                    ps.battery_temp = bt / 10.0

            if os.path.isdir('/sys/class/typec/port0'):
                ps.typec = True
                ps.typec_revision = self.read_sysfs_str('/sys/class/typec/port0/usb_typec_revision')
                ps.typec_pd_revision = self.read_sysfs_str('/sys/class/typec/port0/usb_power_delivery_revision')
                ps.typec_pr = self.read_sysfs_enum('/sys/class/typec/port0/power_role').title()
                ps.typec_dr = self.read_sysfs_enum('/sys/class/typec/port0/data_role').title()
                opmode = self.read_sysfs_str('/sys/class/typec/port0/power_operation_mode')
                ps.typec_operation_mode = self.opmode[opmode]

            if platform:
                ps = platform.process(ps)

            if self.read_sysfs_str(battery_status) == "Unknown":
                if 'usb' in supplies and usb_status:
                    ps.battery_charging = self.read_sysfs_str(usb_status) == "Charging"
                else:
                    ps.battery_charging = ps.battery_current >= 0
                    if ps.battery_current < 0:
                        ps.battery_current *= -1.0
            else:
                ps.battery_charging = self.read_sysfs_str(battery_status) == "Charging"

            for ebat in supplies['bat'][1:]:
                name = 'External battery'
                if os.path.isfile(os.path.join(ebat, 'device/name')):
                    with open(os.path.join(ebat, 'device/name'), 'r') as handle:
                        name = handle.read().strip()
                elif os.path.isfile(os.path.join(ebat, 'model_name')):
                    with open(os.path.join(ebat, 'model_name'), 'r') as name_handle:
                        name = name_handle.read().strip()
                    if os.path.isfile(os.path.join(ebat, 'manufacturer')):
                        with open(os.path.join(ebat, 'manufacturer'), 'r') as manuf_handle:
                            name = '{} {}'.format(manuf_handle.read().strip(), name)

                names = {
                    'kb151': 'Keyboard case',
                    'ip5xxx-battery': 'Keyboard case'
                }
                if name in names:
                    name = names[name]

                ext = ExternalBattery()
                ext.name = name
                ext.present = True
                ext.health = 'N/A'

                ebat_present = os.path.join(ebat, 'present')
                ebat_online = os.path.join(ebat, 'online')
                ebat_status = os.path.join(ebat, 'status')
                ebat_health = os.path.join(ebat, 'health')
                ebat_capacity = os.path.join(ebat, 'capacity')
                ebat_capacity_level = os.path.join(ebat, 'capacity_level')
                ebat_voltage = os.path.join(ebat, 'voltage_now')
                ebat_current = os.path.join(ebat, 'current_now')
                ebat_temp = os.path.join(ebat, 'temp')

                if os.path.isfile(ebat_present):
                    ext.present = self.read_sysfs_str(ebat_present) == '1'
                elif os.path.isfile(ebat_online):
                    ext.present = self.read_sysfs_str(ebat_online) == '1'
                ext.status = self.read_sysfs_str(ebat_status)
                if os.path.isfile(ebat_health):
                    ext.health = self.read_sysfs_str(ebat_health)
                if os.path.isfile(ebat_capacity):
                    ext.capacity = '{}%'.format(self.read_sysfs(ebat_capacity))
                elif os.path.isfile(ebat_capacity_level):
                    ext.capacity = self.read_sysfs_str(ebat_capacity_level)
                if os.path.isfile(ebat_voltage):
                    ext.voltage = self.read_sysfs(ebat_voltage) / 1000000.0
                if os.path.isfile(ebat_current):
                    ext.current = self.read_sysfs(ebat_current) / 1000000.0
                ps.ext.append(ext)

            GLib.idle_add(self.callback, ps)
            time.sleep(1)

    def read_sysfs(self, path, default=None):
        if path is None:
            return 0
        if not os.path.isfile(path):
            return default
        with open(path) as handle:
            return int(handle.read().strip())

    def read_sysfs_str(self, path):
        if path is None:
            return ""
        with open(path) as handle:
            return handle.read().strip()

    def read_sysfs_enum(self, path):
        if path is None:
            return None
        try:
            with open(path) as handle:
                raw = handle.read().strip()
        except:
            return None
        match = self.regex_enum.findall(raw)
        if len(match) > 0:
            return match[0]
        else:
            return None

    def get_supply_types(self):
        result = {}
        mapping = {
            'Battery': 'bat',
            'Mains': 'ac',
            'USB': 'usb'
        }
        for path in glob.glob('/sys/class/power_supply/*'):
            if not os.path.isfile(os.path.join(path, 'type')):
                print("Ignoring {} because it doesn't have a `type` file".format(path))
                continue
            with open(os.path.join(path, 'type')) as handle:
                type = handle.read().strip()

            if type in mapping:
                type = mapping[type]

            if type not in result:
                result[type] = []
            result[type].append(path)
            result[type] = list(sorted(result[type]))

        for type in result:
            print(f'{type}:')
            for path in result[type]:
                print('   ' + path.replace('/sys/class/power_supply/', ''))
        return result


class PowersupplyApplication(Gtk.Application):
    def __init__(self, application_id, flags):
        Gtk.Application.__init__(self, application_id=application_id, flags=flags)
        self.connect("activate", self.new_window)

    def new_window(self, *args):
        AppWindow(self)


class AppWindow:
    def __init__(self, application):
        self.application = application
        builder = Gtk.Builder()
        builder.add_from_resource('/nl/brixit/powersupply/powersupply.glade')
        builder.connect_signals(Handler(builder))

        window = builder.get_object("main_window")
        window.set_application(self.application)
        window.show_all()

        Gtk.main()


class Handler:
    def __init__(self, builder):
        self.builder = builder
        self.window = builder.get_object('main_window')

        self.main_box = builder.get_object('main_box')

        self.supply_type = builder.get_object('supply_type')
        self.bat_capacity = builder.get_object('bat_capacity')
        self.bat_voltage = builder.get_object('bat_voltage')
        self.bat_status = builder.get_object('bat_status')
        self.bat_health = builder.get_object('bat_health')
        self.bat_temp = builder.get_object('bat_temp')
        self.usb_charger = builder.get_object('usb_charger')
        self.usb_voltage = builder.get_object('usb_voltage')
        self.usb_current = builder.get_object('usb_current')
        self.usb_current_row = builder.get_object('usb_current_row')
        self.usb_current_limit = builder.get_object('usb_current_limit')
        self.usb_type = builder.get_object('usb_type')
        self.usb_type_row = builder.get_object('usb_type_row')
        self.typec = builder.get_object('typec')
        self.typec_revision = builder.get_object('typec_revision')
        self.typec_pd_revision = builder.get_object('typec_pd_revision')
        self.typec_pr = builder.get_object('typec_pr')
        self.typec_dr = builder.get_object('typec_dr')
        self.typec_operation_mode = builder.get_object('typec_operation_mode')

        self.boxes = {}

        thread = DataPoller(self.data_update)
        thread.daemon = True
        thread.start()

    def on_quit(self, *args):
        Gtk.main_quit()

    def do_box(self, name, content):
        if name in self.boxes:
            for label in content:
                self.boxes[name][label].set_text(content[label])
            return

        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL)

        heading = Gtk.Label(label=name, xalign=0.0)
        heading.set_margin_bottom(6)
        box.add(heading)

        listbox = Gtk.ListBox()
        box.add(listbox)
        listbox.get_style_context().add_class('content')

        self.boxes[name] = {}
        for label in content:
            row = Gtk.Box()
            row.set_margin_top(8)
            row.set_margin_bottom(8)
            row.set_margin_start(8)
            row.set_margin_end(8)

            dl = Gtk.Label(label=label, xalign=0.0)
            row.pack_start(dl, False, False, 0)

            vl = Gtk.Label(label=content[label])
            row.pack_end(vl, False, False, 0)
            vl.get_style_context().add_class('dim-label')
            self.boxes[name][label] = vl
            listbox.add(row)
        self.main_box.add(box)
        self.main_box.show_all()

    def data_update(self, result):
        if not isinstance(result, PowerStatus):
            return

        self.supply_type.set_text(result.supply_type)

        if result.battery_present:
            if result.battery_charging:
                bat_status = 'Charging'
            else:
                bat_status = 'Discharging'

            if result.battery_voltage is not None and result.battery_current is not None:
                power = result.battery_voltage * result.battery_current
                bat_status += ' ({:0.2f}W)'.format(power)

        else:
            bat_status = 'Not present'
        self.bat_status.set_text(bat_status)

        self.bat_capacity.set_text('{}%'.format(result.battery_capacity))
        self.bat_voltage.set_text('{:0.2f}V'.format(result.battery_voltage))

        if result.battery_temp:
            self.bat_temp.set_text('{:0.1f}°C'.format(result.battery_temp))
        else:
            self.bat_temp.set_text('')
        self.bat_health.set_text(result.battery_health)

        if result.usb_present:
            if result.usb_current_limit is None:
                self.usb_current_limit.set_text("")
            else:
                self.usb_current_limit.set_text('{:0.1f}A'.format(result.usb_current_limit))
            if result.usb_voltage is None:
                self.usb_voltage.set_text('')
            else:
                self.usb_voltage.set_text('{:0.2f}V'.format(result.usb_voltage))
            if result.usb_current is None:
                self.usb_current.set_text('')
                self.usb_current_row.hide()
            else:
                self.usb_current_row.show()
                self.usb_current.set_text('{:0.2f}A'.format(result.usb_current))
            if result.usb_type is None:
                self.usb_type.set_text('')
                self.usb_type_row.hide()
            else:
                self.usb_type_row.show()
                self.usb_type.set_text(result.usb_type)

            self.usb_charger.set_text('Connected')

        elif result.ac_present:
            self.usb_charger.set_text('Connected')
            self.usb_current_limit.set_text('')
            self.usb_voltage.set_text('')
            self.usb_current_row.hide()
            self.usb_type_row.hide()
        else:
            self.usb_current_limit.set_text('')
            self.usb_voltage.set_text('')
            self.usb_charger.set_text('Not connected')
            self.usb_current.set_text('')
            self.usb_current_row.hide()
            self.usb_type_row.hide()

        if result.typec:
            self.typec_revision.set_text(result.typec_revision)
            self.typec_pd_revision.set_text(result.typec_pd_revision)
            self.typec_pr.set_text(result.typec_pr)
            self.typec_dr.set_text(result.typec_dr)
            self.typec_operation_mode.set_text(result.typec_operation_mode)
            self.typec.show_all()
        else:
            self.typec.hide()

        for ext in result.ext:
            content = {
                'Capacity': ext.capacity,
            }
            if ext.voltage is not None:
                content['Voltage'] = f'{ext.voltage}V'

            if ext.current is not None and ext.voltage is not None:
                w = abs(ext.voltage * ext.current)
                content['Status'] = f'{ext.status} ({w:.2f}W)'
            else:
                content['Status'] = ext.status

            content['Health'] = ext.health
            self.do_box(ext.name, content)


def main(version):
    # To make sure libhandy is actually loaded in python
    Handy.init()

    app = PowersupplyApplication("nl.brixit.powersupply", Gio.ApplicationFlags.FLAGS_NONE)
    app.run()


if __name__ == '__main__':
    main('dev')
