# -*- coding: utf-8 -*- {{{
# ===----------------------------------------------------------------------===
#
#                 Installable Component of Eclipse VOLTTRON
#
# ===----------------------------------------------------------------------===
#
# Copyright 2022 Battelle Memorial Institute
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy
# of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
#
# ===----------------------------------------------------------------------===
# }}}

import gevent
import logging
import traceback

from collections import defaultdict
from typing import Any, cast
from weakref import WeakSet

from importlib.metadata import distribution, PackageNotFoundError
try:
    distribution('volttron-core')
    from volttron.client.vip.agent import Agent
except PackageNotFoundError:
    from volttron.platform.vip.agent import Agent

from .interfaces import BaseInterface
from .config import PointConfig, RemoteConfig
from .utils import publication_headers, publish_wrapper
# from platform_driver.poll_scheduler import PollSet  # TODO: This should not import from driver. Need to relocate PollSet.

_log = logging.getLogger(__name__)


class DriverAgent:
    def __init__(self, config: RemoteConfig, core, equipment_model, scalability_test, tz: str, unique_id: Any,
                 vip: Agent.Subsystems):
        self.core = core
        self.equipment_model = equipment_model  # TODO: This should probably move out of the agent and into the base or a library.
        self.scalability_test = scalability_test  # TODO: If this is used from here, it should probably be in the base driver.
        self.tz: str = tz  # TODO: This won't update here if it is updated in the agent. Can it's use be moved out of here?
        self.unique_id: Any = unique_id
        self.vip: Agent.Subsystems = vip
        
        # State Variables
        self.equipment = WeakSet()
        self.heart_beat_value = 0
        self.interface = None
        self.publishers = {}

        try:
            klass = BaseInterface.get_interface_subclass(config.driver_type)
            interface = klass(config, self)
            self.interface = cast(BaseInterface, interface)
        except ValueError as e:
            _log.error(f"Failed to setup device: {e}")
            raise e

        self.config: RemoteConfig = self.interface.config

    def add_registers(self, registry_config: list[PointConfig], base_topic: str):
        """
        Configure a set of registers on this remote.

        :param registry_config: A list of registry points represented as PointConfigs
        :param base_topic: The portion of the topic shared by all points in this registry.
        """
        #_log.debug(f'IN ADD REGISTERS WITH REGISTRY_CONFIG: {registry_config}')
        for register_config in registry_config:
            register = self.interface.create_register(register_config)
            self.interface.insert_register(register, base_topic)
        # TODO: Finalize method is only used by bacnet, currently, and that pauses 30s on each device if it can't connect.
        try:
            self.interface.finalize_setup(initial_setup=True)
        except BaseException as e:
            _log.warning(f'Exception occurred while finalizing setup of interface for {self.unique_id}: {e}.')

        for point_name in self.interface.get_register_names():
            self.update_metadata(point_name)

    def update_metadata(self, point_name):
        register = self.interface.get_register_by_name(point_name)
        point = self.equipment_model.get_node(point_name)
        # TODO: Move the mapping logic hereto an abstract method in the Interface or Register classes?
        if register.register_type == 'bit' or register.python_type is bool:
            ts_type = 'boolean'
        elif register.python_type is int:
            ts_type = 'integer'
        elif register.python_type is float:
            ts_type = 'float'
        elif register.python_type is str:
            ts_type = 'string'
        else:
            ts_type = str(register.python_type)  # TODO: Should this just be "undefined" or something?
        point.meta_data = {
            'units': register.get_units(),
            'type': ts_type,
            'tz': self.tz
        }

    def poll_data(self, poll_set): # PollSet):
        if self.scalability_test:  # TODO: Update scalability testing.
            self.scalability_test.poll_starting(self.unique_id)
        try:
            if not set(poll_set.points.keys()):
                return False
            results, errors = self.interface.get_multiple_points(poll_set.points.keys())
            for failed_point, failure_message in errors.items():
                _log.warning(f'Failed to poll {failed_point}: {failure_message}')
            if results:
                for topic, value in results.items():
                    point = poll_set.points.get(topic)
                    if point and self.equipment_model.is_active(topic):
                        point.last_value = value
                self.publish_poll(results, poll_set)
            return True  # TODO: There could really be better logic in the method to measure success.
        except (Exception, gevent.Timeout) as e:
            _log.error(f'Exception while polling {self.unique_id}: {e}')
            if self.config.debug:
                # TODO: Add an RPC to turn on debugging for individual remotes. Maybe for nodes as well?
                tb = traceback.format_exc()
                _log.error(tb)
            return False
        finally:
            if self.scalability_test:
                self.scalability_test.poll_ending(self.unique_id)

    # noinspection DuplicatedCode
    def publish_poll(self, results, poll_set):
        headers = publication_headers()
        for point_topic in poll_set.single_depth:
            if point_topic in results:
                try:
                    publish_wrapper(self.vip, point_topic, headers=headers, message=[
                        results[point_topic], self.equipment_model.get_node(point_topic).meta_data
                    ])
                except Exception as e:
                    _log.warning(f'Failed to publish single_depth point_topic: {point_topic} -- {e}')
        for point_topic, publish_topic in poll_set.single_breadth:
            if point_topic in results:
                try:
                    publish_wrapper(self.vip, publish_topic, headers=headers, message=[
                        results[point_topic], self.equipment_model.get_node(point_topic).meta_data
                    ])
                except Exception as e:
                    _log.warning(f'Failed to publish single_breadth publish_topic: {publish_topic}'
                                 f' for identifier: {point_topic} -- {e}')
        for device_topic, points in poll_set.multi_depth.items():
            try:
                publish_wrapper(self.vip, f'{device_topic}/multi', headers=headers, message=[
                    {point.rsplit('/', 1)[-1]: results[point] for point in points if point in results},
                    {point.rsplit('/', 1)[-1]: self.equipment_model.get_node(point).meta_data for point in points}
                ])
            except Exception as e:
                _log.warning(f'Failed to publish multi_depth device_topic: {device_topic} -- {e}')
        for publish_topic, points in poll_set.multi_breadth.items():
            try:
                publish_wrapper(self.vip, f'{publish_topic}/multi', headers=headers, message=[
                    {point.rsplit('/', 1)[-1]: results[point] for point in points if point in results},
                    {point.rsplit('/', 1)[-1]: self.equipment_model.get_node(point).meta_data for point in points}
                ])
            except Exception as e:
                _log.warning(f'Failed to publish multi_breadth publish_topic: {publish_topic} -- {e}')

    def heart_beat(self):
        if self.config.heart_beat_point is None:
            return
        self.heart_beat_value = int(not bool(self.heart_beat_value))
        # TODO: config.heart_beat_point should be a set.
        self.set_point(self.config.heart_beat_point, self.heart_beat_value)

    def get_point(self, topic, **kwargs):
        return self.interface.get_point(topic, **kwargs)

    def set_point(self, topic, value, **kwargs):
        return self.interface.set_point(topic, value, **kwargs)

    def get_multiple_points(self, topics, **kwargs):
        return self.interface.get_multiple_points(topics, **kwargs)

    def set_multiple_points(self, topics_values, **kwargs):
        return self.interface.set_multiple_points(topics_values, **kwargs)

    def revert_point(self, topic, **kwargs):
        self.interface.revert_point(topic, **kwargs)

    def revert_all(self, **kwargs):
        self.interface.revert_all(**kwargs)

    def call(self, method_name, topics: list[str], *args, **kwargs):
        if not method_name in self.interface.interface_callable_methods or not hasattr(self.interface, method_name):
            err = (f'Called method {method_name} not found on interface "{self.interface.__class__.__name__}".'
                   f'Check if an expected plugin is not installed.')
            _log.warning(err)
            return {}, {topic: err for topic in topics}
        elif method_name in self.interface.excluded_from_callable_methods:
            err = f'Called method {method_name}, but it excluded from callable methods by configuration.'
            _log.warning(err)
            return {}, {topic: err for topic in topics}
        else:
            return getattr(self.interface, method_name)(self.interface, topics=topics, *args, **kwargs)

    def publish_push(self, results):
        et = self.equipment_model
        headers = publication_headers()
        multi_depth_values, multi_depth_meta, multi_breadth_values, multi_breadth_meta = (
            defaultdict(dict), defaultdict(dict), defaultdict(dict), defaultdict(dict))
        for point_topic, value in results.items():
            point_depth_topic, point_breadth_topic = et.get_point_topics(point_topic)
            device_depth_topic, device_breadth_topic = et.get_device_topics(point_topic)
            point_node = self.equipment_model.get_node(point_topic)
            if point_node and self.equipment_model.is_active(point_topic):
                point_node.last_value = value
            if et.is_published_single_depth(point_topic):
                publish_wrapper(self.vip, point_depth_topic, headers, [value, point_node.meta_data])
            if et.is_published_single_breadth(point_topic):
                publish_wrapper(self.vip, point_breadth_topic, headers, [value, point_node.meta_data])
            if et.is_published_multi_depth(point_topic):
                point_name = point_topic.rsplit('/', 1)[-1]
                multi_depth_values[device_depth_topic][point_name] = value
                multi_depth_meta[device_depth_topic][point_name] = point_node.meta_data
            if et.is_published_multi_breadth(point_topic):
                point_name = point_topic.rsplit('/', 1)[-1]
                multi_breadth_values[device_breadth_topic][point_name] = value
                multi_breadth_meta[device_breadth_topic][point_name] = point_node.meta_data
        if multi_depth_values:
            for device_topic in multi_depth_values:
                publish_wrapper(self.vip, f'{device_topic}/multi', headers,
                                      [multi_depth_values[device_topic], multi_depth_meta[device_topic]])
        if multi_breadth_values:
            for device_topic in multi_breadth_values:
                publish_wrapper(self.vip, f'{device_topic}/multi', headers,
                                [multi_breadth_values[device_topic], multi_breadth_meta[device_topic]])

    def add_equipment(self, device_node):
        # TODO: Is logic needed for scheduling or any other purpose on adding equipment to this remote?
        #_log.debug(f'IN ADD EQUIPMENT, with device_node: {device_node.identifier}')
        self.add_registers([p.config for p in self.equipment_model.points(device_node.identifier)],
                           device_node.identifier)
        self.equipment.add(device_node)

    @property
    def point_set(self):
            return {point for equip in self.equipment for point in self.equipment_model.points(equip.identifier)}
