//go:build e2e_testing
// +build e2e_testing

package router

import (
	"fmt"
	"net"
	"reflect"
	"strconv"
	"sync"

	"github.com/slackhq/nebula"
	"github.com/slackhq/nebula/header"
	"github.com/slackhq/nebula/udp"
)

type R struct {
	// Simple map of the ip:port registered on a control to the control
	// Basically a router, right?
	controls map[string]*nebula.Control

	// A map for inbound packets for a control that doesn't know about this address
	inNat map[string]*nebula.Control

	// A last used map, if an inbound packet hit the inNat map then
	// all return packets should use the same last used inbound address for the outbound sender
	// map[from address + ":" + to address] => ip:port to rewrite in the udp packet to receiver
	outNat map[string]net.UDPAddr

	// All interactions are locked to help serialize behavior
	sync.Mutex
}

type ExitType int

const (
	// Keeps routing, the function will get called again on the next packet
	KeepRouting ExitType = 0
	// Does not route this packet and exits immediately
	ExitNow ExitType = 1
	// Routes this packet and exits immediately afterwards
	RouteAndExit ExitType = 2
)

type ExitFunc func(packet *udp.Packet, receiver *nebula.Control) ExitType

func NewR(controls ...*nebula.Control) *R {
	r := &R{
		controls: make(map[string]*nebula.Control),
		inNat:    make(map[string]*nebula.Control),
		outNat:   make(map[string]net.UDPAddr),
	}

	for _, c := range controls {
		addr := c.GetUDPAddr()
		if _, ok := r.controls[addr]; ok {
			panic("Duplicate listen address: " + addr)
		}
		r.controls[addr] = c
	}

	return r
}

// AddRoute will place the nebula controller at the ip and port specified.
// It does not look at the addr attached to the instance.
// If a route is used, this will behave like a NAT for the return path.
// Rewriting the source ip:port to what was last sent to from the origin
func (r *R) AddRoute(ip net.IP, port uint16, c *nebula.Control) {
	r.Lock()
	defer r.Unlock()

	inAddr := net.JoinHostPort(ip.String(), fmt.Sprintf("%v", port))
	if _, ok := r.inNat[inAddr]; ok {
		panic("Duplicate listen address inNat: " + inAddr)
	}
	r.inNat[inAddr] = c
}

// OnceFrom will route a single packet from sender then return
// If the router doesn't have the nebula controller for that address, we panic
func (r *R) OnceFrom(sender *nebula.Control) {
	r.RouteExitFunc(sender, func(*udp.Packet, *nebula.Control) ExitType {
		return RouteAndExit
	})
}

// RouteUntilTxTun will route for sender and return when a packet is seen on receivers tun
// If the router doesn't have the nebula controller for that address, we panic
func (r *R) RouteUntilTxTun(sender *nebula.Control, receiver *nebula.Control) []byte {
	tunTx := receiver.GetTunTxChan()
	udpTx := sender.GetUDPTxChan()

	for {
		select {
		// Maybe we already have something on the tun for us
		case b := <-tunTx:
			return b

		// Nope, lets push the sender along
		case p := <-udpTx:
			outAddr := sender.GetUDPAddr()
			r.Lock()
			inAddr := net.JoinHostPort(p.ToIp.String(), fmt.Sprintf("%v", p.ToPort))
			c := r.getControl(outAddr, inAddr, p)
			if c == nil {
				r.Unlock()
				panic("No control for udp tx")
			}

			c.InjectUDPPacket(p)
			r.Unlock()
		}
	}
}

// RouteExitFunc will call the whatDo func with each udp packet from sender.
// whatDo can return:
//   - exitNow: the packet will not be routed and this call will return immediately
//   - routeAndExit: this call will return immediately after routing the last packet from sender
//   - keepRouting: the packet will be routed and whatDo will be called again on the next packet from sender
func (r *R) RouteExitFunc(sender *nebula.Control, whatDo ExitFunc) {
	h := &header.H{}
	for {
		p := sender.GetFromUDP(true)
		r.Lock()
		if err := h.Parse(p.Data); err != nil {
			panic(err)
		}

		outAddr := sender.GetUDPAddr()
		inAddr := net.JoinHostPort(p.ToIp.String(), fmt.Sprintf("%v", p.ToPort))
		receiver := r.getControl(outAddr, inAddr, p)
		if receiver == nil {
			r.Unlock()
			panic("Can't route for host: " + inAddr)
		}

		e := whatDo(p, receiver)
		switch e {
		case ExitNow:
			r.Unlock()
			return

		case RouteAndExit:
			receiver.InjectUDPPacket(p)
			r.Unlock()
			return

		case KeepRouting:
			receiver.InjectUDPPacket(p)

		default:
			panic(fmt.Sprintf("Unknown exitFunc return: %v", e))
		}

		r.Unlock()
	}
}

// RouteUntilAfterMsgType will route for sender until a message type is seen and sent from sender
// If the router doesn't have the nebula controller for that address, we panic
func (r *R) RouteUntilAfterMsgType(sender *nebula.Control, msgType header.MessageType, subType header.MessageSubType) {
	h := &header.H{}
	r.RouteExitFunc(sender, func(p *udp.Packet, r *nebula.Control) ExitType {
		if err := h.Parse(p.Data); err != nil {
			panic(err)
		}
		if h.Type == msgType && h.Subtype == subType {
			return RouteAndExit
		}

		return KeepRouting
	})
}

// RouteForUntilAfterToAddr will route for sender and return only after it sees and sends a packet destined for toAddr
// finish can be any of the exitType values except `keepRouting`, the default value is `routeAndExit`
// If the router doesn't have the nebula controller for that address, we panic
func (r *R) RouteForUntilAfterToAddr(sender *nebula.Control, toAddr *net.UDPAddr, finish ExitType) {
	if finish == KeepRouting {
		finish = RouteAndExit
	}

	r.RouteExitFunc(sender, func(p *udp.Packet, r *nebula.Control) ExitType {
		if p.ToIp.Equal(toAddr.IP) && p.ToPort == uint16(toAddr.Port) {
			return finish
		}

		return KeepRouting
	})
}

// RouteForAllExitFunc will route for every registered controller and calls the whatDo func with each udp packet from
// whatDo can return:
//   - exitNow: the packet will not be routed and this call will return immediately
//   - routeAndExit: this call will return immediately after routing the last packet from sender
//   - keepRouting: the packet will be routed and whatDo will be called again on the next packet from sender
func (r *R) RouteForAllExitFunc(whatDo ExitFunc) {
	sc := make([]reflect.SelectCase, len(r.controls))
	cm := make([]*nebula.Control, len(r.controls))

	i := 0
	for _, c := range r.controls {
		sc[i] = reflect.SelectCase{
			Dir:  reflect.SelectRecv,
			Chan: reflect.ValueOf(c.GetUDPTxChan()),
			Send: reflect.Value{},
		}

		cm[i] = c
		i++
	}

	for {
		x, rx, _ := reflect.Select(sc)
		r.Lock()

		p := rx.Interface().(*udp.Packet)

		outAddr := cm[x].GetUDPAddr()
		inAddr := net.JoinHostPort(p.ToIp.String(), fmt.Sprintf("%v", p.ToPort))
		receiver := r.getControl(outAddr, inAddr, p)
		if receiver == nil {
			r.Unlock()
			panic("Can't route for host: " + inAddr)
		}

		e := whatDo(p, receiver)
		switch e {
		case ExitNow:
			r.Unlock()
			return

		case RouteAndExit:
			receiver.InjectUDPPacket(p)
			r.Unlock()
			return

		case KeepRouting:
			receiver.InjectUDPPacket(p)

		default:
			panic(fmt.Sprintf("Unknown exitFunc return: %v", e))
		}
		r.Unlock()
	}
}

// FlushAll will route for every registered controller, exiting once there are no packets left to route
func (r *R) FlushAll() {
	sc := make([]reflect.SelectCase, len(r.controls))
	cm := make([]*nebula.Control, len(r.controls))

	i := 0
	for _, c := range r.controls {
		sc[i] = reflect.SelectCase{
			Dir:  reflect.SelectRecv,
			Chan: reflect.ValueOf(c.GetUDPTxChan()),
			Send: reflect.Value{},
		}

		cm[i] = c
		i++
	}

	// Add a default case to exit when nothing is left to send
	sc = append(sc, reflect.SelectCase{
		Dir:  reflect.SelectDefault,
		Chan: reflect.Value{},
		Send: reflect.Value{},
	})

	for {
		x, rx, ok := reflect.Select(sc)
		if !ok {
			return
		}
		r.Lock()

		p := rx.Interface().(*udp.Packet)

		outAddr := cm[x].GetUDPAddr()
		inAddr := net.JoinHostPort(p.ToIp.String(), fmt.Sprintf("%v", p.ToPort))
		receiver := r.getControl(outAddr, inAddr, p)
		if receiver == nil {
			r.Unlock()
			panic("Can't route for host: " + inAddr)
		}
		r.Unlock()
	}
}

// getControl performs or seeds NAT translation and returns the control for toAddr, p from fields may change
// This is an internal router function, the caller must hold the lock
func (r *R) getControl(fromAddr, toAddr string, p *udp.Packet) *nebula.Control {
	if newAddr, ok := r.outNat[fromAddr+":"+toAddr]; ok {
		p.FromIp = newAddr.IP
		p.FromPort = uint16(newAddr.Port)
	}

	c, ok := r.inNat[toAddr]
	if ok {
		sHost, sPort, err := net.SplitHostPort(toAddr)
		if err != nil {
			panic(err)
		}

		port, err := strconv.Atoi(sPort)
		if err != nil {
			panic(err)
		}

		r.outNat[c.GetUDPAddr()+":"+fromAddr] = net.UDPAddr{
			IP:   net.ParseIP(sHost),
			Port: port,
		}
		return c
	}

	return r.controls[toAddr]
}