# cherry-hue - mqtt device interface for Philips Hue
# Copyright (C) 2021 Bob Carroll <bob.carroll@alum.rit.edu>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import sys
import asyncio
from contextlib import AsyncExitStack
import logging

import yaml
from dns.asyncresolver import resolve
from asyncio_mqtt import Client
import umsgpack
import aiohttp


class Hue(object):
    """
    A simple Hue client.

    :param config: Hue configuration dictionary
    """

    def __init__(self, config):
        self.bridge = config.get('bridge')
        self.username = config.get('username')
        self.url = f'https://{self.bridge}/api/{self.username}'
        self.session = aiohttp.ClientSession()

    async def close(self):
        """
        Performs clean up actions.
        """
        await self.session.close()

    async def get_state(self):
        """
        Retrieves device state from the bridge.

        :returns: a dictionary of device state
        """
        state = {}

        def filter_light(x):
            return {k: v
                    for k, v in x.items()
                    if k in ['on', 'bri', 'hue', 'sat', 'ct', 'colormode']}

        async with self.session.get(f'{self.url}/lights', ssl=False) as resp:
            state['light'] = {k: filter_light(v['state'])
                              for k, v in (await resp.json()).items()}

        async with self.session.get(f'{self.url}/groups', ssl=False) as resp:
            state['group'] = {k: filter_light(v['action'])
                              for k, v in (await resp.json()).items()
                              if v['type'] in ['Room', 'Zone']}

        return state

    async def set_state(self, type_, id_, state):
        """
        Sets the state of a given light or group.

        :param type_: one of light or group
        :param id_: light or group ID
        :param state: target state dictionary
        """
        if type_ == 'light':
            url = f'{self.url}/lights/{id_}/state'
        elif type_ == 'group':
            url = f'{self.url}/groups/{id_}/action'
        else:
            raise Exception(f'{type_} is not a Hue object type')

        if 'ct_k' in state:
            state['ct'] = int(round(1e6 / state['ct_k']))
            del state['ct_k']

        if 'transition' in state:
            state['transitiontime'] = int(round(state['transition'])) * 10
            del state['transition']

        resp = await self.session.put(url, ssl=False, json=state)
        logging.debug(await resp.text())


async def light_state_handler(bridge, messages):
    """
    Event handler for setting the light state from a single request.

    :param bridge: Hue bridge
    :param messages: mqtt message generator
    """
    async for m in messages:
        try:
            platform, device_class, id_, event, _ = m.topic.split('/')
            data = umsgpack.unpackb(m.payload)

            state = {k: v
                     for k, v in data.items()
                     if k in ['on', 'bri', 'ct', 'ct_k', 'transition']}

            logging.debug(f'Setting {device_class} {id_} state: {state}')
            if device_class in ['light', 'group']:
                await bridge.set_state(device_class, id_, state)
        except Exception as ex:
            logging.error(str(ex))


async def light_attr_handler(bridge, messages):
    """
    Event handler for setting light attributes.

    :param bridge: Hue bridge
    :param messages: mqtt message generator
    """
    async for m in messages:
        try:
            platform, device_class, id_, event, _ = m.topic.split('/')
            logging.debug(f'Setting {device_class} {event} state: {m.payload}')

            if event == 'on':
                state = {'on': bool(int(m.payload))}
                await bridge.set_state('light', id_, state)
            elif event in ['bri', 'ct', 'ct_k']:
                state = {event: int(m.payload)}
                await bridge.set_state('light', id_, state)
        except Exception as ex:
            logging.error(str(ex))


async def colour_hack(dc, id_, old_state, new_state):
    """
    Stupid hack to prevent announcing state changes for minor colour value
    differences. Some Hue bulbs report slight changes to the ct field for
    no apparent reason.

    :param dc: device class name
    :param id_: device ID number
    :param old_state: the previous state of the devices
    :param new_state: the current state of the devices
    :returns: True when state change should be suppressed, False otherwise
    """
    if dc not in ['light', 'group']:
        return False
    elif old_state['on'] != new_state['on']:
        return False
    elif old_state['bri'] != new_state['bri']:
        return False
    elif 'colormode' not in old_state:
        return False
    elif old_state['colormode'] != new_state['colormode']:
        return False
    elif old_state['colormode'] != 'ct':
        return False
    elif not (0 < abs(old_state['ct'] - new_state['ct']) < 5):
        return False

    logging.debug(f'Suppressed state change for {dc} {id_}')
    return True


async def publish_device_state(client, dc, old_state, new_state):
    """
    Compares the state of each devices and publishes changes to mqtt.

    :param client: mqtt client
    :param dc: device class name
    :param old_state: the previous state of the devices
    :param new_state: the current state of the devices
    """
    for id_, s in new_state.items():
        if id_ in old_state:
            if s == old_state[id_]:
                continue
            elif await colour_hack(dc, id_, old_state[id_], s):
                continue

        payload = umsgpack.packb(s)
        logging.debug(f'Announcing {dc} {id_} state: {s}')
        await client.publish(f'hue/{dc}/{id_}/state', payload, retain=True)

    for id_ in (x for x in old_state if x not in new_state):
        logging.debug(f'Deleting stale topic for {dc} {id_}')
        await client.publish(f'hue/{dc}/{id_}/state', '', retain=True)


async def publish_all_changes(client, bridge):
    """
    Refreshes the state of all Hue devices and publishes state changes to mqtt.

    :param client: mqtt client
    :param bridge: Hue bridge
    """
    state = {}

    while True:
        try:
            new = await bridge.get_state()

            for dc, ns in new.items():
                os = state.get(dc, {})
                if os != ns:
                    await publish_device_state(client, dc, os, ns)

            state = new
            await asyncio.sleep(1)
        except Exception as ex:
            logging.error(str(ex))


async def set_up_filters(bridge, client, stack):
    """
    Sets up message filters and handlers.

    :param bridge: Hue bridge
    :param client: mqtt client
    :param stack: context manager stack
    :returns: a set of async tasks
    """
    tasks = set()
    topics = {
        'hue/group/+/state/set': light_state_handler,
        'hue/light/+/bri/set': light_attr_handler,
        'hue/light/+/ct/set': light_attr_handler,
        'hue/light/+/ct_k/set': light_attr_handler,
        'hue/light/+/on/set': light_attr_handler,
        'hue/light/+/state/set': light_state_handler}

    for t, cb in topics.items():
        manager = client.filtered_messages(t)
        messages = await stack.enter_async_context(manager)
        task = asyncio.create_task(cb(bridge, messages))
        tasks.add(task)

    return tasks


async def cancel_tasks(tasks):
    """
    Cancels all running async tasks:

    :param tasks: set of tasks to cancel
    """
    for t in tasks:
        if not t.done():
            t.cancel()

        try:
            await t
        except asyncio.CancelledError:
            pass


async def get_broker(config):
    """
    Gets the mqtt broker address from an SRV record.

    :param config: configuration dictionary
    :returns: the broker address
    """
    broker = config.get('mqtt', {}).get('broker')
    if broker is not None:
        return broker

    answer = await resolve('_mqtt._tcp', 'SRV', search=True)
    return next((x.target.to_text() for x in answer))


async def init(config):
    """
    Initializes the Hue agent.

    :param config: configuration dictionary
    """
    bridge = Hue(config.get('hue', {}))

    async with AsyncExitStack() as stack:
        stack.push_async_callback(bridge.close)

        tasks = set()
        stack.push_async_callback(cancel_tasks, tasks)

        client = Client(await get_broker(config), client_id='cherry-hue')
        await stack.enter_async_context(client)
        logging.info('Connected to mqtt broker')

        tasks.update(await set_up_filters(bridge, client, stack))
        await client.subscribe('hue/#')

        task = asyncio.create_task(publish_all_changes(client, bridge))
        tasks.add(task)

        await asyncio.gather(*tasks)


def main():
    """
    CLI entry point.
    """
    if len(sys.argv) != 2:
        print('USAGE: cherry-hue <config file>')
        sys.exit(1)

    with open(sys.argv[1], 'r') as f:
        config = yaml.safe_load(f)

    log = config.get('log', {})
    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s',
                        level=log.get('level', logging.ERROR))

    try:
        asyncio.run(init(config))
    except KeyboardInterrupt:
        pass