#  Copyright 2020 Zeppelin Bend Pty Ltd
#
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

from __future__ import annotations
import logging
from enum import Enum
from typing import Dict, List

from zepben.evolve.model import Measurement
from zepben.evolve.services.common import BaseService
from zepben.evolve.model import ConnectivityNode
from zepben.evolve.model import EnergySource
from zepben.evolve.services.network.tracing import SetPhases
from pathlib import Path

__all__ = ["connect", "NetworkService"]
logger = logging.getLogger(__name__)
TRACED_NETWORK_FILE = str(Path.home().joinpath(Path("traced.json")))


class ProcessStatus(Enum):
    PROCESSED = 0
    INVALID = 1
    SKIPPED = 2


def connect(terminal: Terminal, connectivity_node: ConnectivityNode):
    """
    Connect a `zepben.evolve.iec61970.base.core.terminal.Terminal`` to a `ConnectivityNode`
    `terminal` The `zepben.evolve.iec61970.base.core.terminal.Terminal` to connect.
    `connectivity_node` The `ConnectivityNode` to connect ``zepben.evolve.iec61970.base.core.terminal.Terminal` to.
    """
    terminal.connect(connectivity_node)
    connectivity_node.add_terminal(terminal)


def _attempt_to_reuse_connection(terminal1: Terminal, terminal2: Terminal) -> ProcessStatus:
    """
    Attempt to connect two `zepben.evolve.iec61970.base.core.terminal.Terminal`s.
    Returns `ProcessStatus` reflecting whether the connection was reused. PROCESSED if a connection was
    established, INVALID if it couldn't be, and SKIPPED if neither terminal had an existing `ConnectivityNode`.
    """
    cn1 = terminal1.connectivity_node
    cn2 = terminal2.connectivity_node

    if cn1 is not None:
        if cn2 is not None:
            if cn1 is cn2:
                return ProcessStatus.PROCESSED
            elif connect(terminal2, cn1.mrid):
                return ProcessStatus.PROCESSED
            return ProcessStatus.INVALID
    elif cn2 is not None:
        return ProcessStatus.PROCESSED if connect(terminal1, cn2.mrid) else ProcessStatus.INVALID
    return ProcessStatus.SKIPPED


class NetworkService(BaseService):
    """
    A full representation of the power network.
    Contains a map of equipment (string ID's -> Equipment/Nodes/etc)
    **All** `IdentifiedObject's` submitted to this Network **MUST** have unique mRID's!

    Attributes -
        metrics_store : Storage for meter measurement data associated with this network.
    """

    name: str = "network"
    _connectivity_nodes: Dict[str, ConnectivityNode] = dict()
    _auto_cn_index: int = 0
    _measurements: Dict[str, List[Measurement]] = []

    def __init__(self):
        self._objectsByType[ConnectivityNode] = self._connectivity_nodes

    def get_measurements(self, mrid: str, t: type) -> List[Measurement]:
        """
        Get all measurements of type `t` associated with the given `mrid`.
                                                                                                              
        The `mrid` should be either a `zepben.evolve.iec61970.base.core.power_system_resource.PowerSystemResource` or a
        `zepben.evolve.iec61970.base.core.terminal.Terminal` MRID that is assigned to the corresponding fields on the measurements.
        Returns all `Measurement`s indexed by `mrid` in this service.
        Raises `KeyError` if `mrid` isn't present in this service.
        """
        return [meas for meas in self._measurements[mrid] if isinstance(meas, t)]

    def add_measurement(self, measurement: Measurement) -> bool:
        """
        Add a `zepben.evolve.cim.iec61970.base.meas.measurement.Measurement` to this `NetworkService`

        `measurement` The `Measurement` to add.
        Returns `True` if `measurement` was added, `False` otherwise
        """
        return self._index_measurement(measurement) and self.add(measurement)

    def remove_measurement(self, measurement) -> bool:
        """
        Remove a `zepben.evolve.cim.iec61970.base.meas.measurement.Measurement` from this `NetworkService`

        `measurement` The `Measurement` to remove.
        Returns `True` if `measurement` was removed, `False` otherwise
        """
        self._remove_measurement_index(measurement)
        return self.remove(measurement)

    def connect_by_mrid(self, terminal: Terminal, connectivity_node_mrid: str) -> bool:
        """
        Connect a `zepben.evolve.iec61970.base.core.terminal.Terminal` to the `ConnectivityNode` with mRID `connectivity_node_mrid`
        `terminal` The `zepben.evolve.iec61970.base.core.terminal.Terminal` to connect.
        `connectivity_node_mrid` The mRID of the `ConnectivityNode`. Will be created in the `Network` if it
        doesn't already exist.
        Returns True if the connection was made or already existed, False if `zepben.evolve.iec61970.base.core.terminal.Terminal` was already connected to a
        different `ConnectivityNode`
        """
        if not connectivity_node_mrid:
            return False

        if terminal.connectivity_node:
            return connectivity_node_mrid == terminal.connectivity_node.mrid

        cn = self.add_connectivitynode(connectivity_node_mrid)
        connect(terminal, cn)
        return True

    def connect_terminals(self, terminal1: Terminal, terminal2: Terminal) -> bool:
        """
        Connect two `zepben.evolve.iec61970.base.core.terminal.Terminal`s
        Returns True if the `zepben.evolve.iec61970.base.core.terminal.Terminal`s could be connected, False otherwise.
        """
        status = _attempt_to_reuse_connection(terminal1, terminal2)
        if status == ProcessStatus.PROCESSED:
            return True
        elif status == ProcessStatus.INVALID:
            return False

        cn = self.add_connectivitynode(self._generate_cn_mrid())
        connect(terminal2, cn)
        connect(terminal1, cn)

        return True

    def _generate_cn_mrid(self):
        mrid = f"generated_cn_{self._auto_cn_index}"
        while mrid in self._connectivity_nodes:
            self._auto_cn_index += 1
            mrid = f"generated_cn_{self._auto_cn_index}"
        return mrid

    def disconnect(self, terminal: Terminal):
        """
        Disconnect a `zepben.evolve.iec61970.base.core.terminal.Terminal`` from its `ConnectivityNode`. Will also remove the `ConnectivityNode` from this
        `Network` if it no longer has any terminals.
        `terminal` The `zepben.evolve.iec61970.base.core.terminal.Terminal` to disconnect.
        """
        cn = terminal.connectivity_node
        if cn is None:
            return
        cn.remove_terminal(terminal)
        terminal.disconnect()
        if cn.num_terminals() == 0:
            del self._connectivity_nodes[cn.mrid]

    def disconnect_by_mrid(self, connectivity_node_mrid: str):
        """
        Disconnect a `ConnectivityNode` from this `Network`. Will disconnect all ``zepben.evolve.iec61970.base.core.terminal.Terminal`s from the
        `ConnectivityNode`
        `connectivity_node_mrid` The mRID of the `ConnectivityNode` to disconnect.
        Raises `KeyError` if there is no `ConnectivityNode` for `connectivity_node_mrid`
        """
        cn = self._connectivity_nodes[connectivity_node_mrid]
        if cn is not None:
            for term in cn.terminals:
                term.disconnect()
            cn.clear_terminals()
            del self._connectivity_nodes[connectivity_node_mrid]

    def get_primary_sources(self):
        """
        Get the primary source for this network. All directions are applied relative to this EnergySource
        Returns The primary EnergySource
        """
        return [source for source in self._objectsByType[EnergySource].values() if source.has_phases()]

    def add_connectivitynode(self, mrid: str):
        """
        Add a connectivity node to the network.
        `mrid` mRID of the ConnectivityNode
        Returns A new ConnectivityNode with `mrid` if it doesn't already exist, otherwise the existing
                 ConnectivityNode represented by `mrid`
        """
        if mrid not in self._connectivity_nodes:
            self._connectivity_nodes[mrid] = ConnectivityNode(mrid=mrid)
            return self._connectivity_nodes[mrid]
        else:
            return self._connectivity_nodes[mrid]

    async def set_phases(self):
        set_phases = SetPhases()
        await set_phases.run(self)

    def _index_measurement(self, measurement: Measurement, mrid: str) -> bool:
        if not mrid:
            return False

        if mrid in self._measurements:
            for meas in self._measurements[mrid]:
                if meas.mrid == measurement.mrid:
                    return False
            else:
                self._measurements[mrid].append(measurement)
                return True
        else:
            self._measurements[mrid] = [measurement]
            return True

    def _remove_measurement_index(self, measurement: Measurement):
        try:
            self._measurements[measurement.terminal_mrid].remove(measurement)
        except KeyError:
            pass
        try:
            self._measurements[measurement.power_system_resource_mrid].remove(measurement)
        except KeyError:
            pass




