from datetime import datetime, timedelta
import json
from logging import Logger
import uuid
from pytradfri import Gateway
from pytradfri.api.libcoap_api import APIFactory
from pytradfri.group import Group
from pytradfri.command import Command
from pytradfri.error import RequestTimeout
from pytradfri.const import (
    ATTR_DEVICE_STATE,
    ATTR_LIGHT_COLOR_HEX,
    ATTR_LIGHT_DIMMER
)
from typing import Any, Callable, Iterable, Union

import numpy

CONFIG_FILE = "tradfri_psk.conf"

class RGB(numpy.ndarray):
  @classmethod
  def from_str(cls, hex: str):
    return numpy.array([int(hex[i:i+2], 16) for i in (0, 2, 4)]).view(cls)

  def __str__(self):
    self = self.astype(numpy.uint8)
    return ''.join(format(n, 'x') for n in self)

class TradfriHandler:
    def __init__(self, gateway_hostname: str, key: str, logger: Logger):
        self.logger = logger
        conf = self.load_psk(CONFIG_FILE)

        try:
            identity = conf[gateway_hostname].get("identity")
            psk = conf[gateway_hostname].get("key")
            api_factory = APIFactory(host=gateway_hostname, psk_id=identity, psk=psk)
        except KeyError:
            identity = uuid.uuid4().hex
            api_factory = APIFactory(host=gateway_hostname, psk_id=identity)

        psk = api_factory.generate_psk(key)

        conf[gateway_hostname] = {"identity": identity, "key": psk}
        self.save_psk(CONFIG_FILE, conf)
        self.api = api_factory.request
        self.gateway = Gateway()
        self.load_groups()

    @staticmethod
    def load_psk(filename: str) -> dict:
        try:
            with open(filename, encoding="utf-8") as fdesc:
                return json.loads(fdesc.read())
        except FileNotFoundError:
            return {}

    @staticmethod
    def save_psk(filename: str, config: dict):
        data = json.dumps(config, sort_keys=True, indent=4)
        with open(filename, "w", encoding="utf-8") as fdesc:
            fdesc.write(data)

    @staticmethod
    def average_hex_color(colors: list[str]):
        if len(colors) == 1:
            return colors[0]
        rgb_colors = [RGB.from_str(hex) for hex in colors]
        return (numpy.sum(rgb_colors, axis=0) // len(rgb_colors)).view(RGB)

    def export_group(self, group: Group) -> dict[str, Any]:
        # These properties exists on the group as well, but they are incorrect for some reason
        hex_colors, states = zip(*map(lambda light: (light.light_control.lights[0].hex_color,
                                                     light.light_control.lights[0].state),
                                      filter(lambda device: device.has_light_control,
                                             self.group_members[group.id])
                                      ))
        return {
            "name": group.name,
            "id": group.id,
            "state": any(states),
            "dimmer": group.dimmer,
            "color": '#' + str(self.average_hex_color(list(hex_colors)))
        }

    def export_groups(self) -> list[str]:
        return list(map(self.export_group, self.get_groups()))
    
    def load_group_members(self, group: Group):
        try:
            self.group_members[group.id] = self.api(group.members())
        except RequestTimeout:
            self.logger.error("Trådfri timed out!")

    def load_groups(self):
        self.groups: dict[int, Group] = {}
        self.group_members: dict[int, Iterable] = {}
        try:
            devices_commands = self.api(self.gateway.get_groups())
            groups = self.api(devices_commands)
            for group in groups:
                self.groups[group.id] = group
                self.load_group_members(group)
            self.groups_last_updated = datetime.now()
        except RequestTimeout:
            self.logger.error("Trådfri timed out!")
            self.groups_last_updated = None

    def get_groups(self) -> Iterable[Group]:
        # Only update every 5 minutes at most
        if (len(self.groups) == 0 or not self.groups_last_updated or
            datetime.now() > self.groups_last_updated + timedelta(minutes=5)):
            self.load_groups()
        return self.groups.values()

    def set_state(self, group_id: int, new_state: bool) -> bool:
        return self.run_api_command_for_group(lambda lg: lg.set_state(new_state),
                                              lambda lg: self.update_group(lg, ATTR_DEVICE_STATE, int(new_state)),
                                              group_id)

    def set_dimmer(self, group_id: int, value: int) -> bool:
        return self.run_api_command_for_group(lambda lg: lg.set_dimmer(value, transition_time=1),
                                              lambda lg: self.update_group(lg, ATTR_LIGHT_DIMMER, value),
                                              group_id)

    def set_hex_color(self, group_id: int, value: str) -> bool:
        return self.run_api_command_for_group(lambda lg: lg.set_hex_color(value, transition_time=1),
                                              lambda lg: self.update_group(lg, ATTR_LIGHT_COLOR_HEX, value),
                                              group_id)

    def run_api_command_for_group(self,
                                  command_function: Callable[[Group], Command],
                                  update_function: Callable[[Group], None],
                                  group_id: int) -> bool:
        light_group = self.groups[group_id]
        try:
            self.api(command_function(light_group))
        except RequestTimeout:
            return False
        update_function(light_group)
        return True

    # This is a bit hacky, but allows to update the state of the device without
    # refetching it through the gateway
    def update_group(self, group: Group, key: str, new_value: Union[int, str]):
        group.raw[key] = new_value

        for member in self.group_members[group.id]:
            if not member.has_light_control:
                continue
            for light in member.light_control.lights:
                light.raw[key] = new_value

@@ -404,13 +404,8 @@
 port
- Any,
  Cal
@@ -2278,11 +2278,23 @@
 tr, 
-Any
+Union%5Bstr, int%5D
 %5D:%0A 

Remove Any from type hints