from argparse import ArgumentParser
from collections import defaultdict, OrderedDict
from itertools import combinations
import math
import random
import sys

import simpy
import yaml
import networkx as nx

from .core import TickProcess

from .communities import types as communities
from .routers import types as routers
from .traces import types as traces

class Network:
    ''''''
    def __init__(self, env,
                 packets=None,
                 node_factory=None,
                 community=None,
                 trace=None):
        ''''''
        self.env = env

        # contact trace
        if trace is None:
            trace = traces['random']()
        self.trace = trace
        self.trace.start(env, self)

        # community detection
        self.community = community
        if community is not None:
            self.community.start(env, self)

        # packet generation
        if packets is None:
            packets = {}
        self.packets = PacketGenerator(**packets)
        self.packets.start(env, self)

        # create node network
        if node_factory is None:
            node_factory = NodeFactory(tick_rate=1, router=routers['direct'])
        self.nodes = [
            node_factory(env, self, nid)
            for nid in range(self.trace.nodes)
        ]
        self.links = [
            (a, b)
            for a, b in combinations(self.nodes, 2)
        ]

        # set up networkx graph
        self.graph = nx.Graph()
        self.graph.add_nodes_from(self.nodes)
        self.graph.add_edges_from([
            (a, b, { 'state': False })
            for a, b in self.links
        ])

    def set_link(self, a, b, state):
        if isinstance(a, int):
            a = self.nodes[a]

        if isinstance(b, int):
            b = self.nodes[b]

        edge = self[a][b]
        if edge['state'] == state:
            return

        if state is None:
            state = not edge['state']
        edge['state'] = state

        if self.community:
            self.community.set_link(a, b, state, self.env.now)

    def toggle_link(self, a, b):
        self.set_link(a, b, None)

    def send_link(self, a, b, packet):
        ''''''
        if self.graph[a][b]['state']:
            # TODO: transfer delay
            b.recv(packet)
        else:
            raise Exception('Nodes {} and {} not connected'.format(a, b))

    def __getitem__(self, node):
        ''''''
        return self.graph[node]


def NodeFactory(router, **kwargs):
    nid = -1
    def factory(env, network, nid):
        nid += 1
        return Node(env, network, nid, router=router, **kwargs)
    return factory


class Node(TickProcess):
    ''''''
    def __init__(self, env, network, nid,
                 buffer_size=None, tick_time=1, router=None):
        ''''''
        super().__init__(tick_time)
        self.env = env

        self.network = network
        self.id = nid

        self.buffer = Buffer(self.env, capacity=buffer_size)

        # bind router as a class method
        if router is None:
            router = routers['direct']
        self.router = router #router.__get__(self, Node)
        self.router_state = {}

        self.start(env)

    def route_packets(self):
        packets_to_delete = []

        for packet in self.buffer:
            if self.router(self, packet, self.router_state):
                packets_to_delete.append(packet)

        for packet in packets_to_delete:
            self.buffer.remove(packet)

    def process(self):
        ''''''
        while True:
            self.buffer.clean()
            self.route_packets()
            yield self.tick()

    def send(self, to, packet, reason=None):
        # TODO: transfer delay
        packet.send(self, to, reason=reason)
        self.network.send_link(self, to, packet)

    def recv(self, packet):
        if packet.destination == self:
            packet.recv()
        else:
            self.buffer.add(packet)

    '''
    def __eq__(self, other):
        return self.id == other.id
    '''

    @property
    def community(self):
        return self.network.community[self]

    @property
    def links(self):
        '''
        Returns a list of connected links.
        '''
        links = {
            met: data
            for met, data in self.network[self].items()
            if data['state']
        }
        return links

    def __repr__(self):
        return 'Node(id={})'.format(self.id)


class Buffer:
    def __init__(self, env, capacity=0):
        self.env = env


        if capacity is None or capacity <= 0:
            self.capacity = float('inf')
        else:
            self.capacity = capacity

        self.buffer = OrderedDict()
        self.used = 0

    def clean(self):
        packets_to_drop = []

        for packet in self:
            if packet.ttl < self.env.now:
                packets_to_drop.append(packet)

        for packet in packets_to_drop:
            self.remove(packet)

    def add(self, packet):
        if self.used < self.capacity:
            self.used += 1
            self.buffer[packet] = None
        else:
            raise Exception('buffer full')

    def remove(self, packet):
        self.used -= 1
        del self.buffer[packet]

    def __contains__(self, packet):
        return packet in self.buffer

    def __iter__(self):
        return iter(self.buffer)

    def __len__(self):
        return len(self.buffer)

class PacketGenerator(TickProcess):
    def __init__(self, ticks_per_packet=1, start_delay=0, time_to_live=60):
        super().__init__(ticks_per_packet)

        if time_to_live is None:
            raise ValueError('time_to_live must be specified')

        self.start_delay = start_delay
        self.time_to_live = time_to_live
        self.packets = []

    def process(self, network):
        packet_id = 0
        def create(packet_id):
            source, dest = random.choice(network.links)
            packet = Packet(packet_id, source, dest,
                            self.env.now + self.time_to_live, None)
            self.packets.append(packet)
            source.recv(packet)
        yield self.env.timeout(self.start_delay)
        print('starting packets')
        #@DEBUG
        for i in range(217):
            pass
        #@DEBUG
        while True:
            create(packet_id)
            yield self.tick()
            packet_id += 1

    @property
    def recieved(self):
        return len([
            packet
            for packet in self.packets
            if packet.recieved
        ])

    @property
    def recieved_count(self):
        return sum([
            packet.recieved_count
            for packet in self.packets
        ])

    @property
    def sent(self):
        return sum([
            packet.sent
            for packet in self.packets
        ])

    @property
    def stats(self):
        stats = {}
        for packet in self.packets:
            for stat, value in packet.stats.items():
                if stat not in stats:
                    stats[stat] = value
                else:
                    stats[stat] += value
        return stats

    @property
    def paths(self):
        return {
            packet: packet.path
            for packet in self.packets
            if packet.recieved
        }

    @property
    def total(self):
        return len(self.packets)

    def __str__(self):
        ret = "recieved: {}/{}, delivery ratio: {}, packets: {}, sent: {}, stats: {}, path: {}".format(
            self.recieved,
            self.recieved_count,
            self.recieved / len(self.packets),
            len(self.packets),
            self.sent,
            self.stats,
            sum([ len(path) for path in self.paths.values() ]) / self.recieved
        )
        '''
        for packet, path in self.paths.items():
            ret += '\n  {}: {}'.format(packet, path)
        '''
        return ret

class Packet:
    def __init__(self, id, source, destination, ttl, payload):
        self.id = id
        self.source = source
        self.destination = destination
        self.ttl = ttl
        self.payload = payload

        self.path = []

        self.stats = dict()

        self.sent = 0

        self.recieved = False
        self.recieved_count = 0

        self.dropped = False
        self.dropped_count = 0

    def send(self, a, b, reason=None):
        if reason is None:
            self.path.append((a.id, b.id))
        else:
            self.path.append((a.id, b.id, reason))
        self.sent += 1

    def drop(self):
        self.dropped = True
        self.dropped_count += 1

    def recv(self):
        self.recieved = True
        self.recieved_count += 1

    def __str__(self):
        return "Packet(id={}, src={}, dst={})".format(
            self.id,
            self.source,
            self.destination
        )