from ryu.base import app_manager
from ryu.controller import ofp_event
from ryu.controller.handler import CONFIG_DISPATCHER, MAIN_DISPATCHER
from ryu.controller.handler import set_ev_cls
from ryu.ofproto import ofproto_v1_3
from ryu.lib import dpid as dpid_lib
from ryu.lib import stplib
from ryu.lib.packet import packet
from ryu.lib.packet import ethernet
from ryu.lib.packet import vlan
from ryu.lib.packet import ether_types

import util
topo = util.load_src("topo_query", "topo/topo_query.py").TopoQuery("topo/topo3.json")

class CherryPick(app_manager.RyuApp):
	OFP_VERSIONS = [ofproto_v1_3.OFP_VERSION]
	_CONTEXTS = {'stplib': stplib.Stp}

	def __init__(self, *args, **kwargs):
		super(CherryPick, self).__init__(*args, **kwargs)
		# define MAC address table
		self.mac_to_port = {}
		self.stp = kwargs['stplib']

	@set_ev_cls(ofp_event.EventOFPSwitchFeatures, CONFIG_DISPATCHER)
	def switch_features_handler(self, ev):
		# datapath == the OpenFlow switch that issued this message
		datapath = ev.msg.datapath
		ofproto = datapath.ofproto
		parser = datapath.ofproto_parser

		# default table-miss flow entry
		actions = [parser.OFPActionOutput(ofproto.OFPP_CONTROLLER,
                                          ofproto.OFPCML_NO_BUFFER)]
		self.add_flow(datapath, 0, parser.OFPMatch(), actions)

	def add_flow(self, datapath, priority, match, actions, buffer_id=None):
		ofproto = datapath.ofproto
		parser = datapath.ofproto_parser

		# immediately use the specified action
		inst = [parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS,
                                             actions)]
		if buffer_id:
			mod = parser.OFPFlowMod(datapath=datapath, buffer_id=buffer_id,
                                    priority=priority, match=match,
                                    instructions=inst)
		else:
			mod = parser.OFPFlowMod(datapath=datapath, priority=priority,
                                    match=match, instructions=inst)
		datapath.send_msg(mod)

	def delete_flow(self, datapath):
		ofproto = datapath.ofproto
		parser = datapath.ofproto_parser

		for dst in self.mac_to_port[datapath.id].keys():
			match = parser.OFPMatch(eth_dst=dst)
			mod = parser.OFPFlowMod(
				datapath, command=ofproto.OFPFC_DELETE,
				out_port=ofproto.OFPP_ANY, out_group=ofproto.OFPG_ANY,
				priority=1, match=match)
			datapath.send_msg(mod)

	def get_vlan_actions(self, dpid, in_port, src_mac, dst_mac, vids, parser):
		# not in any pods
		if topo.get_pod_id(dpid) == None or \
			topo.get_pod_id(src_mac) == None or \
			topo.get_pod_id(dst_mac) == None:
			return None
		pick = False
		if topo.get_switch_type(dpid) == "core":
			pick = True
		elif topo.get_switch_type(dpid) == "agg":
			pre = topo.get_switch_type(topo.get_next_hop_id(dpid, in_port))
			if pre == "toc" \
				and topo.get_pod_id(dpid) == topo.get_pod_id(dst_mac):
				pick = True
		elif topo.get_switch_type(dpid) == "toc":
			pre_id = topo.get_next_hop_id(dpid, in_port)
			if not pre_id == "host":
				pre = topo.get_switch_type(dpid)
				if pre == "agg" \
					and topo.get_pod_id(dpid) == topo.get_pod_id(src_mac)\
					and topo.get_pod_id(dst_mac) != topo.get_pod_id(src_mac):
					pick = True
		if pick:
			link_id = int(topo.get_link_id(dpid, in_port))
			self.logger.info("push vlan: %s", link_id)
			act1 = [parser.OFPActionPushVlan(),
				    parser.OFPActionSetField(vlan_vid=(0x1000 | link_id<<6))]
			vid = 0x1000
			if len(vids) > 0:
				vid = vid | vids[0]
			act2 = [parser.OFPActionSetField(vlan_vid=(vid | link_id))]
			return [act1, act2]
		return None

	@set_ev_cls(stplib.EventPacketIn, MAIN_DISPATCHER)
	def _packet_in_handler(self, ev):
		# If you hit this you might want to increase
		# the "miss_send_length" of your switch
		if ev.msg.msg_len < ev.msg.total_len:
			self.logger.debug("packet truncated: only %s of %s bytes",
                              ev.msg.msg_len, ev.msg.total_len)
		msg = ev.msg
		datapath = msg.datapath
		ofproto = datapath.ofproto
		parser = datapath.ofproto_parser
		in_port = msg.match['in_port']

		pkt = packet.Packet(msg.data)
		eth = pkt.get_protocols(ethernet.ethernet)[0]

		if eth.ethertype == ether_types.ETH_TYPE_LLDP:
			# ignore lldp packet
			return
		dst = eth.dst
		src = eth.src

		# vlan test
		vls = pkt.get_protocols(vlan.vlan)
		vids = list(map(lambda e: e.vid, vls))
		self.logger.info("vids %s", vids)

		# switch id
		dpid = datapath.id
		self.mac_to_port.setdefault(dpid, {})

		self.logger.info("\033[32mpacket in\033[0m [%s:%s] %s --> %s", dpid, in_port, src, dst)

		# learn a mac address to avoid FLOOD next time.
		self.mac_to_port[dpid][src] = in_port

		if dst in self.mac_to_port[dpid]:
			out_port = self.mac_to_port[dpid][dst]
		else:
			out_port = ofproto.OFPP_FLOOD

		#
		vlan_act = False
		actions = None
		tagged = len(vids) > 0
		if out_port != ofproto.OFPP_FLOOD:
			actions = self.get_vlan_actions(dpid, in_port, src, dst, vids, parser)
		if actions == None:
			actions = [parser.OFPActionOutput(out_port)]
		else:
			vlan_act = True
			actions[0].append(parser.OFPActionOutput(out_port))
			actions[1].append(parser.OFPActionOutput(out_port))

		# install a flow to avoid packet_in next time
		if out_port != ofproto.OFPP_FLOOD:
			match = parser.OFPMatch(in_port=in_port, eth_dst=dst)
			notag = parser.OFPMatch(in_port=in_port, eth_dst=dst,
									vlan_vid=0x0000)
			hastag = parser.OFPMatch(in_port=in_port, eth_dst=dst,
									 vlan_vid=(0x1000, 0x1000))
			# remove vlan tag before sending to host
			if topo.get_next_hop_id(dpid, out_port) == "host":
				self.logger.info("\033[32mto host\033[0m %s:%s", dpid, out_port)
				# assert vlan_act == False
				to_con = [parser.OFPActionOutput(ofproto.OFPP_CONTROLLER,
                                          		 ofproto.OFPCML_NO_BUFFER)]
				self.add_flow(datapath, 30, match, to_con)
			else:
				if not vlan_act:
					if msg.buffer_id != ofproto.OFP_NO_BUFFER:
						self.add_flow(datapath, 20, match, actions, msg.buffer_id)
						return
					else:
						self.add_flow(datapath, 20, match, actions)
				else:
					if msg.buffer_id != ofproto.OFP_NO_BUFFER:
						self.add_flow(datapath, 20, notag, actions[0], msg.buffer_id)
						self.add_flow(datapath, 20, hastag, actions[1], msg.buffer_id)
						return
					else:
						self.add_flow(datapath, 20, notag, actions[0])
						self.add_flow(datapath, 20, hastag, actions[1])

		data = None
		if msg.buffer_id == ofproto.OFP_NO_BUFFER:
			data = msg.data

		if vlan_act:
			if not tagged:
				out = parser.OFPPacketOut(datapath=datapath, buffer_id=msg.buffer_id,
										  in_port=in_port, actions=actions[0], data=data)
			else:
				out = parser.OFPPacketOut(datapath=datapath, buffer_id=msg.buffer_id,
										  in_port=in_port, actions=actions[1], data=data)
		else:
			if tagged:
				actions.insert(0, parser.OFPActionPopVlan())
			out = parser.OFPPacketOut(datapath=datapath, buffer_id=msg.buffer_id,
									  in_port=in_port, actions=actions, data=data)
		datapath.send_msg(out)

	@set_ev_cls(stplib.EventTopologyChange, MAIN_DISPATCHER)
	def _topology_change_handler(self, ev):
		dp = ev.dp
		dpid_str = dpid_lib.dpid_to_str(dp.id)
		msg = 'Receive topology change event. Flush MAC table.'
		self.logger.debug("\033[32m[dpid=%s]\033[0m %s", dpid_str, msg)

		if dp.id in self.mac_to_port:
			self.delete_flow(dp)
			del self.mac_to_port[dp.id]

	@set_ev_cls(stplib.EventPortStateChange, MAIN_DISPATCHER)
	def _port_state_change_handler(self, ev):
		dpid_str = dpid_lib.dpid_to_str(ev.dp.id)
		of_state = {stplib.PORT_STATE_DISABLE: 'DISABLE',
					stplib.PORT_STATE_BLOCK: 'BLOCK',
					stplib.PORT_STATE_LISTEN: 'LISTEN',
					stplib.PORT_STATE_LEARN: 'LEARN',
					stplib.PORT_STATE_FORWARD: 'FORWARD'}
		self.logger.debug("\033[32m[dpid=\033[0m%s\033[32m][port=\033[0m%d\033[32m] state=\033[0m%s",
						  dpid_str, ev.port_no, of_state[ev.port_state])