package net

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

type LuxRoute struct {
	Key         crypto.LuxKey
	Destination *net.UDPAddr
	Associated  *LuxChannel
	Nonces      LuxNonceList
}

// !!!!!
// TODO: map key can be destination ID, while Key (+ ID) in route struct would be SOURCE
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) 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
	}

	r.routes[key.Id] = &LuxRoute{
		Key:         key,
		Destination: channel.Address,
		Associated:  r.addOutboundChannel(channel),
		Nonces:      NewLuxNonceList(),
	}
	return nil
}

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

	r.routes[r.thisKey.Id] = &LuxRoute{
		Key:         r.thisKey,
		Destination: channel.Address,
		Associated:  r.addInboundChannel(channel),
		Nonces:      NewLuxNonceList(),
	}
	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.Key.Id]; !ok {
		return
	}

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

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

	var err error
	var packet LuxPacket

	/* BUG:
	 *	nodes share same UUID. When updating routes by UUID, it will overwrite others' nodes routes.
	 *	this is current limitation of protocol
	 */

	// first we look key from routes
	if route, ok := r.GetRoute(dgram.Target); ok {
		packet, err = DecryptLuxPacket(dgram, route.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.Key.Id.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.Key.Id.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.Key.Id.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 {
					r.routes[key.Id] = &LuxRoute{
						Key:         key,
						Destination: dgram.Target,
						Associated:  dgram.Channel,
						Nonces:      NewLuxNonceList(),
					}
					route = r.routes[key.Id]

					log.Debugf("established route %s <-> %s", route.Key.Id.String(), route.Destination.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")
	}

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

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

func (r *LuxRouter) Multicast(packet LuxPacket, group proto.LuxType) error {
	for _, route := range r.routes {
		if route.Key.Type == group {
			packet.Nonce = GenerateLuxNonce()
			dgram, err := EncryptLuxPacket(packet, route.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
}