package net

import (
	"bytes"
	"errors"
	"fmt"
	"lux/crypto"
	"lux/proto"
	"net"
	"sync"
)

type LuxRouteType int

const (
	LuxRouteFromSource = 0
	LuxRouteToTarget   = 1
)

type LuxRoute struct {
	Type        LuxRouteType
	Target      proto.LuxID
	Source      proto.LuxID
	Destination *net.UDPAddr
	Associated  *LuxChannel
	Nonces      LuxNonceList
}

func (route *LuxRoute) String() string {
	var dir string
	switch route.Type {
	case LuxRouteFromSource:
		dir = "  <-> *"
	case LuxRouteToTarget:
		dir = "* <->  "
	default:
		dir = "  <?>  "
	}

	return fmt.Sprintf("%s%s%s %s %s",
		route.Target.String(), dir, route.Source.String(),
		route.Associated.Type.String(), route.Destination.IP.String())
}

/*
Routing Map Table

	key TargetID <-> value Route (SourceID, SourceKey, Transport..)

Host routing:

	node 1111 <-> source 0001 host A
	node 1112 <-> source 0001 host A
	node 1113 <-> source 0001 host A
	node 1114 <-> source 0001 host A

Node routing:

	host 0001 <-> source 1111 node A
	host 0002 <-> source 1111 node A
	host 0003 <-> source 1111 node A
	host 0004 <-> source 1111 node A
	node 1112 <-> source 1111 node A

Node and host:

	Receives and decrypts with host key, which is TargetID (as well as LuxPacket Target)
	Sends to host with host key, which is TargetID

Node 1111 and node 1112:

	Node 1111 receives and decrypts with node 1112 key, which is TargetID
	Node 1111 sends to node 1112 with node 1112 key, which is TargetID

But host uses host key when communicating to node, so target ID in host routing table cannot
be used. To overcome this, a "direction" field must be introduced to routing entry, that
decides whether use target or source ID-key.
*/
type LuxRouter struct {
	thisKey  crypto.LuxKey
	keyStore crypto.LuxKeyStore

	routes map[proto.LuxID]*LuxRoute

	channelLock sync.RWMutex
	outbound    []LuxChannel
	inbound     []LuxChannel

	dgramChan chan LuxDatagram
}

func NewLuxRouter(key crypto.LuxKey, ks crypto.LuxKeyStore) LuxRouter {
	return LuxRouter{
		thisKey:   key,
		keyStore:  ks,
		routes:    make(map[proto.LuxID]*LuxRoute, 0),
		outbound:  make([]LuxChannel, 0),
		inbound:   make([]LuxChannel, 0),
		dgramChan: make(chan LuxDatagram),
	}
}

func (r *LuxRouter) GetThisKey() crypto.LuxKey {
	return r.thisKey
}

func (r *LuxRouter) GetRouterType() proto.LuxType {
	return r.thisKey.Type
}

func (r *LuxRouter) addOutboundChannel(ch LuxChannel) *LuxChannel {
	r.channelLock.Lock()

	r.outbound = append(r.outbound, ch)
	channel := &r.outbound[len(r.outbound)-1]

	r.channelLock.Unlock()
	return channel
}

func (r *LuxRouter) addInboundChannel(ch LuxChannel) *LuxChannel {
	r.channelLock.Lock()

	r.inbound = append(r.inbound, ch)
	channel := &r.inbound[len(r.inbound)-1]

	r.channelLock.Unlock()
	return channel
}

// the ID is not destination, but rather peer associated for this route, like source ID.
// Destination router always know who is he, therefore we dont need target ID.
func (r *LuxRouter) CreateOutboundRoute(id proto.LuxID, chType LuxChannelType, udpAddr string) error {
	// we gonna look up key by id from key store
	key, ok := r.keyStore.Get(id)
	if !ok {
		return errors.New("key not found")
	}

	// create outbound channel
	channel, err := NewLuxOutboundChannel(udpAddr, chType)
	if err != nil {
		return err
	}

	var routeType LuxRouteType
	switch r.GetRouterType() {
	case proto.LuxTypeHost:
		if key.Type == proto.LuxTypeNode {
			routeType = LuxRouteFromSource
		} else {
			routeType = LuxRouteToTarget
		}
	case proto.LuxTypeNode:
		routeType = LuxRouteToTarget
	}

	r.routes[key.Id] = &LuxRoute{
		Type:        routeType,
		Target:      id,
		Source:      r.thisKey.Id,
		Destination: channel.Address,
		Associated:  r.addOutboundChannel(channel),
		Nonces:      NewLuxNonceList(),
	}
	log.Debugf("outbound route: %s", r.routes[key.Id].String())

	return nil
}

func (r *LuxRouter) CreateInboundChannel(chType LuxChannelType, udpAddr string) error {
	channel, err := NewLuxInboundChannel(udpAddr, chType)
	if err != nil {
		return err
	}

	r.addInboundChannel(channel)
	return nil
}

// close channel when error happened
func (r *LuxRouter) CloseChannel(channel *LuxChannel, closeInbound bool) {
	r.channelLock.Lock()
	for i, ch := range r.outbound {
		if &ch == channel {
			r.outbound = append(r.outbound[:i], r.outbound[i+1:]...)
		}
	}

	if closeInbound {
		for i, ch := range r.inbound {
			if &ch == channel {
				r.inbound = append(r.inbound[:i], r.inbound[:i+1]...)
			}
		}
	}
	r.channelLock.Unlock()
}

func (r *LuxRouter) GetDgramChannel() chan<- LuxDatagram {
	return r.dgramChan
}

// goroutine to receive datagrams and send them to router over channel
func channelReceiver(r *LuxRouter, channel *LuxChannel) {
	dgramChan := r.GetDgramChannel()

	var dgram LuxDatagram
	var err error

	for err == nil {
		dgram, err = channel.Recv()
		dgramChan <- dgram
	}

	r.CloseChannel(channel, true)
}

func (r *LuxRouter) Start() {
	r.channelLock.RLock()
	for _, inbound := range r.inbound {
		go channelReceiver(r, &inbound)
	}
	r.channelLock.RUnlock()
}

func (r *LuxRouter) Stop() {
	// close all channels
	r.channelLock.Lock()
	for _, inbound := range r.inbound {
		inbound.Close()
	}
	r.inbound = r.inbound[0:]

	for _, outbound := range r.outbound {
		outbound.Close()
	}
	r.outbound = r.outbound[0:]
	r.channelLock.Unlock()

	r.routes = make(map[proto.LuxID]*LuxRoute)
}

func (r *LuxRouter) RecvDgram() LuxDatagram {
	return <-r.dgramChan
}

func udpAddrEqual(addr *net.UDPAddr, other *net.UDPAddr) bool {
	return addr.IP.Equal(other.IP) && addr.Port == other.Port
}

func (r *LuxRouter) GetRoutes() map[proto.LuxID]*LuxRoute {
	return r.routes
}

func (r *LuxRouter) GetRoute(udpAddr *net.UDPAddr) (*LuxRoute, bool) {
	for _, route := range r.routes {
		if udpAddrEqual(route.Destination, udpAddr) {
			return route, true
		}
	}

	return nil, false
}

func (r *LuxRouter) DeleteRoute(route *LuxRoute) {
	if _, ok := r.routes[route.Target]; !ok {
		return
	}

	r.CloseChannel(route.Associated, false)
	delete(r.routes, route.Target)
}

func (r *LuxRouter) GetRouteKey(route *LuxRoute) (crypto.LuxKey, bool) {
	switch route.Type {
	case LuxRouteFromSource:
		return r.keyStore.Get(route.Source)
	case LuxRouteToTarget:
		return r.keyStore.Get(route.Target)
	}

	return crypto.LuxKey{}, false
}

func (r *LuxRouter) Recv() (LuxPacket, error) {
	dgram := r.RecvDgram()

	var err error
	var packet LuxPacket

	// first we look key from routes
	if route, ok := r.GetRoute(dgram.Target); ok {
		key, _ := r.GetRouteKey(route)
		packet, err = DecryptLuxPacket(dgram, key)
		if err != nil {
			// do we really fail here?
			log.Debugf("DecryptLuxPacket err %v for route %v", err, route)
			return packet, err
		}

		// check if LuxID matches
		if !bytes.Equal(packet.Target.UUID[:], route.Target.UUID[:]) {
			// not matches.. we discard route and throw away packet
			log.Infof("packet from %s received at route %v mismatches associated target UUID %s",
				dgram.Target.String(), route, route.Target.String())

			r.DeleteRoute(route)
			return packet, errors.New("bogus packet from established route")
			/* NOTE:
			 *	there may be rare situation where multiple clients behind NAT/CGNAT
			 *	reusing same IP:Port of established route, confusing routing table.
			 *	This is not critical, since hosts and nodes are expected to send
			 *	data again after interval of time, and packet loss is expected and
			 *	tolerated at protocol design level.
			 */
		}

		// packet arrived to right route and successfully decrypted
		return packet, nil
	} else {
		// first time seeing peer - bruteforce keys from keystore
		for _, key := range r.keyStore.Keys() {
			packet, err = DecryptLuxPacket(dgram, key)
			if err != nil {
				continue
			}

			if bytes.Equal(packet.Target.UUID[:], key.Id.UUID[:]) {
				// key UUID and decrypted UUID matching - create OR update route and return packet
				var ok bool
				var route *LuxRoute
				if route, ok = r.routes[packet.Target]; ok {
					log.Debugf("updating route %s: %s -> %s", route.Target.String,
						route.Destination.String(), dgram.Target.String())

					route.Destination = dgram.Target
					route.Associated = dgram.Channel
					// since packet arrived from different transport, we flush nonces
					route.Nonces.Flush()
				} else {
					var routeType LuxRouteType
					if key.Type == proto.LuxTypeHost {
						routeType = LuxRouteFromSource
					} else {
						routeType = LuxRouteToTarget
					}

					r.routes[key.Id] = &LuxRoute{
						Type:        routeType,
						Target:      key.Id,
						Source:      r.thisKey.Id,
						Destination: dgram.Target,
						Associated:  dgram.Channel,
						Nonces:      NewLuxNonceList(),
					}
					route = r.routes[key.Id]

					log.Debugf("established route: %s", route.String())
				}

				// rotate nonce
				if !route.Nonces.RotateOrFail(packet.Nonce) {
					// failed nonce, discard packet
					log.Debug("failed nonce")
					return packet, fmt.Errorf("packet failed nonce check")
				}

				packet.ChannelType = dgram.Channel.Type
				return packet, nil
			}
		}
	}

	return packet, fmt.Errorf("non-peer packet from %s", dgram.Target.String())
}

func (r *LuxRouter) Send(packet LuxPacket) error {
	route, ok := r.routes[packet.Target]
	if !ok {
		return errors.New("no route to peer")
	}
	key, _ := r.GetRouteKey(route)

	packet.Nonce = GenerateLuxNonce()
	dgram, err := EncryptLuxPacket(packet, key, route.Destination)
	if err != nil {
		return err
	}

	// TODO: close route if it fails?
	return route.Associated.Send(dgram)
}

// In multicast we send packet to all routes by group. Group is lux peer type
// which is always determined by target ID-key, even in LuxRouteFromSource,
// since target is node, but we encrypt with source ID-key.
func (r *LuxRouter) Multicast(packet LuxPacket, group proto.LuxType) error {
	for _, route := range r.routes {
		key, _ := r.GetRouteKey(route)
		targetKey, _ := r.keyStore.Get(route.Target)

		if targetKey.Type == group {
			packet.Nonce = GenerateLuxNonce()
			dgram, err := EncryptLuxPacket(packet, key, route.Destination)
			if err != nil {
				return err
			}

			if err = route.Associated.Send(dgram); err != nil {
				// TODO: close route if it fails?
				return err
			}
		}
	}

	return nil
}