package net

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

type LuxRouteType int

const (
	LuxRouteHostToNode = 0
	LuxRouteNodeToNode = 1
)

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

// !!!!!
// TODO: map key can be destination ID, while Key (+ ID) in route struct would be SOURCE
/*
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 channelTypeToRouteType(chType LuxChannelType) LuxRouteType {
	switch chType {
	case LuxChannelInterior:
		return LuxRouteNodeToNode
	case LuxChannelExterior:
		return LuxRouteHostToNode
	default:
		log.Panicf("can't translate chType %d to lux route type", chType)
	}
}

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{
		Type:        channelTypeToRouteType(chType),
		Target:      id,
		Source:      r.thisKey.Id,
		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{
		Source:      r.thisKey.Id,
		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.Target]; !ok {
		return
	}

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

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 {
		key, _ := r.keyStore.Get(route.Target)
		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 {
					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
}