package node

import (
	"log"
	"net"
	"net/netip"
	"sync/atomic"
	"vppn/fasttime"
)

// TODO:
type connRouter interface {
	Lookup(byte) *peer
	Mediator() *peer
}

type connWriter struct {
	*net.UDPConn
	localIP  byte
	buf      []byte
	counters [256]uint64
	lookup   func(byte) *peer
}

func newConnWriter(conn *net.UDPConn, localIP byte, lookup func(byte) *peer) *connWriter {
	w := &connWriter{
		UDPConn: conn,
		localIP: localIP,
		buf:     make([]byte, bufferSize),
		lookup:  lookup,
	}

	for i := range w.counters {
		w.counters[i] = uint64(fasttime.Now() << 30)
	}

	return w
}

func (w *connWriter) WriteTo(remoteIP, packetType byte, data []byte) error {
	peer := w.lookup(remoteIP)
	if peer == nil || peer.Addr == nil {
		log.Printf("No peer: %d", remoteIP)
		return nil
	}

	h := header{
		Counter:    atomic.AddUint64(&w.counters[remoteIP], 1),
		SourceIP:   w.localIP,
		ViaIP:      0,
		DestIP:     remoteIP,
		PacketType: packetType,
	}

	buf := encryptPacket(&h, peer.SharedKey, data, w.buf)

	_, err := w.WriteToUDPAddrPort(buf, *peer.Addr)
	return err
}

// ----------------------------------------------------------------------------

type connReader struct {
	*net.UDPConn
	localIP   byte
	dupChecks [256]*dupCheck
	lookup    func(byte) *peer
	buf       []byte
}

func newConnReader(conn *net.UDPConn, localIP byte, lookup func(byte) *peer) *connReader {
	r := &connReader{
		UDPConn: conn,
		localIP: localIP,
		lookup:  lookup,
		buf:     make([]byte, bufferSize),
	}
	for i := range r.dupChecks {
		r.dupChecks[i] = newDupCheck(0)
	}
	return r
}

func (r *connReader) Read(buf []byte) (remoteAddr netip.AddrPort, h header, data []byte, err error) {
	var n int

	for {
		n, remoteAddr, err = r.ReadFromUDPAddrPort(buf[:bufferSize])
		if err != nil {
			return
		}

		data = buf[:n]

		if n < headerSize {
			continue // Packet it soo short.
		}

		h.Parse(data)
		if len(data) != headerSize+int(h.DataSize) {
			log.Printf("Incorrect size")
			continue // Packet is corrupt.
		}

		peer := r.lookup(h.SourceIP)
		if peer == nil {
			log.Printf("No peer...")
			continue
		}

		out, ok := decryptPacket(peer.SharedKey, data, r.buf)
		if !ok {
			log.Printf("Decrypt failed...")
			continue
		}

		out, data = data, out

		if r.dupChecks[h.SourceIP].IsDup(h.Counter) {
			log.Printf("Duplicate...")
			continue
		}

		return
	}
}