from .engineobj import SqPandasEngine
from dataclasses import dataclass
import os

import pandas as pd
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

from suzieq.exceptions import EmptyDataframeError


# TODO:
#  topology for different VRFs?
#  iBGP vs eBGP?
#  color by device type?
#  physical topology without LLDP -- is this possible?
#  how to draw multiple topologies
#  be able to ask if a node has neighbors by type (physical, overlay, protocol, etc)
#  questions
#    * without knowing hierarchy, labels or tags it's unclear how to group things for good picture
# how could we add state of connection (like Established) per protocol

graph_output_dir = '/tmp/suzieq-graphs'


class TopologyObj(SqPandasEngine):

    @staticmethod
    def table_name():
        return 'topology'

    def _init_dfs(self, namespaces):
        """Initialize the dataframes used"""

        self._if_df = self._get_table_sqobj('interfaces') \
            .get(namespace=namespaces, state="up",
                 columns=['namespace', 'hostname', 'ifname', 'ipAddressList',
                          'ip6AddressList', 'state', 'type', 'master', 'macaddr'])

        if self._if_df.empty:
            raise EmptyDataframeError(f"No interface found for {namespaces}")
        self._if_df['vrf'] = self._if_df.apply(
            lambda x: x['master'] if x['type'] not in ['bridge', 'bond_slave']
            else 'default', axis=1)
        self._if_df['vrf'] = np.where(self._if_df['vrf'] == '', 'default',
                                      self._if_df['vrf'])
        self._if_df = self._if_df.explode('ipAddressList') \
                                 .explode('ip6AddressList') \
                                 .fillna({'ipAddressList': '0.0.0.0/0',
                                          'ip6AddressList': '::0/0'})
        self._if_df['plen'] = self._if_df.ipAddressList.str.split('/').str[1]
        self._if_df['ipAddress'] = self._if_df.ipAddressList.str.split(
            '/').str[0]
        self._if_df = self._if_df.drop(
            columns=['ipAddressList', 'ip6AddressList'])

        self.nses = self._if_df['namespace'].unique()

    def get(self, **kwargs):
        self._namespaces = kwargs.get("namespace", self.ctxt.namespace)
        hostname = kwargs.pop('hostname', [])
        user_query = kwargs.pop('query_str', '')
        polled = kwargs.pop('polled', '')
        via = kwargs.pop('via', [])
        ifname = kwargs.pop('ifname', '')
        peerHostname = kwargs.pop('peerHostname', [])
        columns = kwargs.pop('columns', ['default'])

        self._init_dfs(self._namespaces)
        self.lsdb = pd.DataFrame()
        self._a_df = pd.DataFrame()
        self._ip_table = pd.DataFrame()

        self.services = [
            Services('lldp', {}, ['ifname', 'vrf'],
                     self._augment_lldp_show),
            Services('arpnd', {},
                     ['ipAddress', 'macaddr', 'ifname', 'vrf', 'arpndBidir'],
                     self._augment_arpnd_show),
            Services('bgp', {'state': 'Established',
                             'columns': ['*']},
                     ['vrf'],
                     self._augment_bgp_show),
            # Services('evpnVni', {}, 'peerHostname',
            #    'vni', self._augment_evpnvni_show),
            Services('ospf', {'state': 'full'},
                     ['ifname', 'vrf'], self._augment_ospf_show),
        ]

        if via:
            self.services = [x for x in self.services if x.name in via]
        key = 'peerHostname'
        for srv in self.services:
            if 'columns' not in srv.extra_args:
                srv.extra_args['columns'] = columns
            df = self._get_table_sqobj(srv.name).get(
                **kwargs,
                **srv.extra_args
            )

            if not df.empty:
                if srv.augment:
                    df = srv.augment(df)
                cols = ['namespace', 'hostname', key]
                if srv.extra_cols:
                    cols = cols + srv.extra_cols
                df = df[cols]
                df.insert(len(df.columns), srv.name, True)
                if self.lsdb.empty:
                    self.lsdb = df
                else:
                    self.lsdb = self.lsdb.merge(df, how='outer')

        self._find_polled_neighbors(polled)
        if self.lsdb.empty:
            return self.lsdb

        if 'ipAddress' not in self.lsdb.columns:
            self.lsdb['ipAddress'] = ''

        self.lsdb = self.lsdb[~self.lsdb.hostname.isna()] \
                        .drop_duplicates() \
                        .rename({'ipAddress': 'peerIP', 'macaddr': 'peerMac'},
                                axis=1, errors='ignore') \
                        .dropna(subset=['peerHostname', 'peerIP'], how='all') \
                        .fillna({'peerHostname': 'unknown',
                                 'ifname': 'unknown', 'arpnd': False,
                                 'peerIP': '', 'peerMac': '',
                                 'arpndBidir': False,
                                 'bgp': False, 'ospf': False,
                                 'lldp': False, 'vrf': 'N/A'})

        self.lsdb['vrf'] = np.where(self.lsdb.vrf == "bridge", "-",
                                    self.lsdb.vrf)
        self.lsdb['peerHostname'] = np.where(
            self.lsdb.peerHostname == "unknown", self.lsdb.peerIP,
            self.lsdb.peerHostname)
        self.lsdb = self.lsdb.drop(columns=['peerIP', 'peerMac'],
                                   errors='ignore')
        self.lsdb = self.lsdb.query('peerHostname != ""')

        # Apply the appropriate filters
        if not self.lsdb.empty:
            if hostname:
                self.lsdb = self.lsdb.query(f'hostname.isin({hostname})')
            if peerHostname:
                self.lsdb = self.lsdb.query(
                    f'peerHostname.isin({peerHostname})')
            if ifname:
                self.lsdb = self.lsdb.query(f'ifname.isin({ifname})')

        if user_query and not self.lsdb.empty:
            self.lsdb = self.lsdb.query(user_query)

        cols = self.lsdb.columns.tolist()
        if 'timestamp' in self.lsdb.columns:
            cols.remove('timestamp')
            cols.append('timestamp')

        return self.lsdb[cols].reset_index(drop=True)

    def _find_polled_neighbors(self, polled):
        if self.lsdb.empty:
            return

        devices = self._get_table_sqobj('device').get(namespace=self._namespaces,
                                                      columns=['namespace', 'hostname'])
        self.lsdb = devices.merge(self.lsdb, how='outer',
                                  indicator=True)
        self.lsdb = self.lsdb.rename(columns={'_merge': 'polled'})
        self.lsdb.polled = self.lsdb.polled == 'both'

        if polled and polled.lower() == 'true':
            self.lsdb = self.lsdb.query('polled == True')
        elif polled.lower() == 'false':
            self.lsdb = self.lsdb.query('polled == False')
        if self.lsdb.empty:
            self.lsdb = pd.DataFrame()

    def _create_graphs_from_lsdb(self):
        self.graphs = {}
        for ns, df in self.lsdb.groupby(by=['namespace']):
            attrs = [srv.name for srv in self.services if srv.name in df.columns]
            self.graphs[ns] = nx.from_pandas_edgelist(
                df, 'hostname', 'peerHostname', attrs, nx.MultiGraph)

    # TODO: eventually this needs to move to ospf after we figure out the
    #   schema augmentation story
    def _augment_ospf_show(self, df):
        if not df.empty:
            df = df.query('adjState != "passive"').reset_index(drop=True)

        return df

    def _augment_lldp_show(self, df):
        if not df.empty:
            df = df[df.peerHostname != '']
            df = df.merge(
                self._if_df[['namespace', 'hostname', 'ifname', 'vrf']],
                on=['namespace', 'hostname', 'ifname'])
        return df

    def _augment_bgp_show(self, df):
        return df

    def _augment_evpnvni_show(self, df):
        if not df.empty:
            df = df.explode('remoteVtepList').dropna(how='any')
            df = df.rename(columns={'remoteVtepList': 'peerIP'})
            df = df.merge(self.ip_table, on=['namespace', 'peerIP'],
                          how='left').dropna(how='any').drop_duplicates()
        return df

    def _augment_arpnd_show(self, df):

        if not df.empty:
            # weed out entries that are not reachable
            df = df.query('state != "failed"').reset_index(drop=True)
            # Get the VRF
            df = df.merge(
                self._if_df[['namespace', 'hostname', 'ifname', 'master']],
                left_on=['namespace', 'hostname', 'oif'],
                right_on=['namespace', 'hostname', 'ifname'],
                suffixes=['', '_y'], how='outer') \
                .rename(columns={'master': 'vrf'}) \
                .dropna(subset=['ipAddress'])

            df['vrf'] = np.where(df['vrf'] == '', 'default', df['vrf'])
            df = df.drop(columns=['vlan', 'oif', 'mackey', 'remoteVtepIp',
                                  'timestamp_y'], errors='ignore')

            # address are where I find the neighbor, these own the mac
            addr = self.address_df[['namespace', 'hostname', 'macaddr']] \
                .drop_duplicates() \
                .rename(columns={'hostname': 'peerHostname'})
            df = df.merge(addr, on=['namespace', 'macaddr'], how='left') \
                   .dropna(subset=['hostname', 'ipAddress']) \
                   .rename(columns={'oif': 'ifname'}) \
                   .drop_duplicates()

            # Use MAC table entries to find the local port for a MAC on an SVI
            mac_df = self._get_table_sqobj('macs')\
                         .get(namespace=self._namespaces, localOnly=True,
                              columns=['namespace', 'hostname', 'vlan',
                                       'macaddr', 'oif'])
            df = df.merge(mac_df,
                          on=['namespace', 'hostname', 'macaddr'],
                          how='outer') \
                .dropna(subset=['hostname'])
            df['ifname'] = np.where(
                df['oif'].isnull(), df['ifname'], df['oif'])

            df['arpndBidir'] = df.apply(
                lambda x, y: True
                if not y.query(f'namespace=="{x.namespace}" and '
                               f'hostname=="{x.peerHostname}" and '
                               f'peerHostname=="{x.hostname}"').empty
                else False, args=(df,), axis=1)

        self._arpnd_df = df
        return self._arpnd_df

    @property
    def address_df(self):
        if self._a_df.empty:
            self._a_df = self._if_df
        return self._a_df

    @property
    def ip_table(self):
        if self._ip_table.empty:
            addr = self.address_df
            if not addr.empty:
                self._ip_table = addr[['namespace',
                                       'hostname', 'ipAddressList']]
                self._ip_table = self._ip_table.explode(
                    'ipAddressList').dropna(how='any')
                self._ip_table = self._ip_table \
                                     .rename(
                                         columns={'ipAddressList': 'peerIP',
                                                  'hostname': 'peerHostname'})
                self._ip_table['peerIP'] = self._ip_table['peerIP'] \
                                               .str.replace("/.+", "")
                self._ip_table = self._ip_table[self._ip_table['peerIP'] != '-']
        return self._ip_table

    def summarize(self, **kwargs):
        self.get(**kwargs)
        if self.lsdb.empty:
            return self.lsdb

        self.ns = {}
        for i in self.nses:
            self.ns[i] = {}
        self._create_graphs_from_lsdb()
        self._analyze_lsdb_graph()

        return pd.DataFrame(self.ns)

    def _analyze_lsdb_graph(self):
        for ns in self.nses:
            for name in [srv.name for srv in self.services if
                         nx.get_edge_attributes(self.graphs[ns], srv.name)]:

                G = nx.Graph([(s, d, data) for s, d, data in
                              self.graphs[ns].edges(data=True) if data[name] == True])
                if G.nodes:
                    self.ns[ns][f'{name}_number_of_nodes'] = len(G.nodes)
                    self.ns[ns][f'{name}_number_of_edges'] = len(G.edges)
                    if not nx.is_connected(G):
                        self.ns[ns][f'{name}_is_fully_connected'] = False
                        self.ns[ns][f'{name}_center'] = False
                    else:
                        self.ns[ns][f'{name}_is_fully_connected'] = True
                        self.ns[ns][f'{name}_center'] = nx.barycenter(G)

                    self.ns[ns][f'{name}_self_loops'] = list(
                        nx.nodes_with_selfloops(G))

                    self.ns[ns][f'{name}_number_of_disjoint_sets'] = len(
                        list(nx.connected_components(G)))

                    self.ns[ns][f'{name}_degree_histogram'] = nx.degree_histogram(
                        G)

                    # if there are too many degrees than the column gets too big
                    if len(self.ns[ns][f'{name}_degree_histogram']) > 6:
                        self.ns[ns][f'{name}_degree_histogram'] = '...'

                else:
                    for k in [f'{name}_is_fully_connected', f'{name}_center',
                              f'{name}_self_loops', f'{name}_number_of_disjoint_sets',
                              f'{name}_degree_histogram', f'{name}_number_of_nodes',
                              f'{name}_number_of_edges']:
                        self.ns[ns][k] = None

    def _make_images(self):
        if not os.path.exists(graph_output_dir):
            os.makedirs(graph_output_dir)
        for ns in self.nses:
            pos = nx.spring_layout(self.graphs[ns])
            for name in [srv.name for srv in self.services if
                         nx.get_edge_attributes(self.graphs[ns], srv.name)]:
                edges = [(u, v) for u, v, d in self.graphs[ns].edges(data=True) if
                         d[name] == True]
                if len(edges) > 1:

                    nx.draw_networkx_nodes(
                        self.graphs[ns], pos=pos, node_size=25)
                    if len(self.graphs[ns].nodes) < 20:
                        nx.draw_networkx_labels(
                            self.graphs[ns], pos=pos, font_size=8)

                    nx.draw_networkx_edges(
                        self.graphs[ns], edgelist=edges, pos=pos)
                    plt.savefig(f"{graph_output_dir}/{ns}_{name}.png")
                    print(f"created {graph_output_dir}/{ns}_{name}.png")
                    plt.close()

    def unique(self, **kwargs):
        '''Unique values for topology'''
        columns = kwargs.pop('columns', ['default'])
        count = kwargs.pop("count", 0)
        query_str = kwargs.get('query_str', '')

        df = self.get(columns=['*'], **kwargs)
        column = columns[0]

        if not count:
            return (pd.DataFrame({f'{column}': df[column].unique()}))
        else:
            r = df[column].value_counts()
            return (pd.DataFrame({column: r})
                    .reset_index()
                    .rename(columns={column: 'numRows',
                                     'index': column})
                    .sort_values(column))


@dataclass(frozen=True)
class Services:
    name: str
    extra_args: dict
    extra_cols: list
    augment: any
