package tunnel

import (
	"context"
	"encoding/gob"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/inverse-inc/packetfence/go/chisel/share/cio"
	"github.com/inverse-inc/packetfence/go/chisel/share/settings"
	"github.com/jpillora/sizestr"
	"golang.org/x/crypto/ssh"
	"golang.org/x/sync/errgroup"
)

// listenUDP is a special listener which forwards packets via
// the bound ssh connection. tricky part is multiplexing lots of
// udp clients through the entry node. each will listen on its
// own source-port for a response:
//
//	                                            (random)
//	src-1 1111->...                         dst-1 6345->7777
//	src-2 2222->... <---> udp <---> udp <-> dst-1 7543->7777
//	src-3 3333->...    listener    handler  dst-1 1444->7777
//
// we must store these mappings (1111-6345, etc) in memory for a length
// of time, so that when the exit node receives a response on 6345, it
// knows to return it to 1111.
func listenUDP(l *cio.Logger, sshTun sshTunnel, remote *settings.Remote) (*udpListener, error) {
	a, err := net.ResolveUDPAddr("udp", remote.Local())
	if err != nil {
		return nil, l.Errorf("resolve: %s", err)
	}
	conn, err := net.ListenUDP("udp", a)
	if err != nil {
		return nil, l.Errorf("listen: %s", err)
	}
	//ready
	u := &udpListener{
		Logger:  l,
		sshTun:  sshTun,
		remote:  remote,
		inbound: conn,
	}
	return u, nil
}

type udpListener struct {
	*cio.Logger
	sshTun      sshTunnel
	remote      *settings.Remote
	inbound     *net.UDPConn
	outboundMut sync.Mutex
	outbound    *udpChannel
	sent, recv  int64
}

func (u *udpListener) run(ctx context.Context) error {
	defer u.inbound.Close()
	//udp doesnt accept connections,
	//udp simply forwards packets
	//and therefore only needs to listen
	eg, ctx := errgroup.WithContext(ctx)
	ctx, cancel := context.WithCancel(ctx)
	eg.Go(func() error {
		return u.runInbound(ctx)
	})
	eg.Go(func() error {
		return u.runOutbound(ctx)
	})
	if u.remote.Dynamic {
		eg.Go(func() error {
			return u.monitorInactivity(ctx, cancel)
		})
	}
	if err := eg.Wait(); err != nil {
		u.Debugf("listen: %s", err)
		return err
	}
	u.Debugf("Close (sent %s received %s)", sizestr.ToString(u.sent), sizestr.ToString(u.recv))
	return nil
}

func (u *udpListener) runInbound(ctx context.Context) error {
	const maxMTU = 9012
	buff := make([]byte, maxMTU)
	for !isDone(ctx) {
		//read from inbound udp
		u.inbound.SetReadDeadline(time.Now().Add(time.Second))
		n, addr, err := u.inbound.ReadFromUDP(buff)
		if e, ok := err.(net.Error); ok && (e.Timeout() || e.Temporary()) {
			continue
		}
		if ctx.Err() != nil {
			u.Debugf("Reading on a canceled remote")
			return nil
		}
		if err != nil {
			return u.Errorf("read error: %w", err)
		}
		//upsert ssh channel
		uc, err := u.getUDPChan(ctx)
		if err != nil {
			if strings.HasSuffix(err.Error(), "EOF") {
				continue
			}
			return u.Errorf("inbound-udpchan: %w", err)
		}
		//send over channel, including source address
		b := buff[:n]
		if err := uc.encode(addr.String(), b); err != nil {
			if strings.HasSuffix(err.Error(), "EOF") {
				u.Infof("dropped packet")
				continue //dropped packet...
			}
			return u.Errorf("encode error: %w", err)
		}
		//stats
		atomic.AddInt64(&u.sent, int64(n))
	}
	return nil
}

func (u *udpListener) runOutbound(ctx context.Context) error {
	for !isDone(ctx) {
		//upsert ssh channel
		uc, err := u.getUDPChan(ctx)
		if err != nil {
			if strings.HasSuffix(err.Error(), "EOF") {
				continue
			}
			return u.Errorf("outbound-udpchan: %w", err)
		}
		//receive from channel, including source address
		p := udpPacket{}
		if err := uc.decode(&p); err == io.EOF {
			u.Infof("ssh disconnected,")
			//outbound ssh disconnected, get new connection...
			continue
		} else if err != nil {
			return u.Errorf("decode error: %w", err)
		}
		//write back to inbound udp
		addr, err := net.ResolveUDPAddr("udp", p.Src)
		if err != nil {
			return u.Errorf("resolve error: %w", err)
		}
		n, err := u.inbound.WriteToUDP(p.Payload, addr)
		if err != nil {
			return u.Errorf("write error: %w", err)
		}
		//stats
		atomic.AddInt64(&u.recv, int64(n))
	}
	return nil
}

func (u *udpListener) getUDPChan(ctx context.Context) (*udpChannel, error) {
	u.outboundMut.Lock()
	defer u.outboundMut.Unlock()
	//cached
	if u.outbound != nil {
		return u.outbound, nil
	}
	//not cached, bind
	sshConn := u.sshTun.getSSH(ctx)
	if sshConn == nil {
		return nil, fmt.Errorf("ssh-conn nil")
	}
	//ssh request for udp packets for this proxy's remote,
	//just "udp" since the remote address is sent with each packet
	dstAddr := u.remote.Remote() + "/udp|" + u.remote.Handler
	rwc, reqs, err := sshConn.OpenChannel("chisel", []byte(dstAddr))
	if err != nil {
		return nil, fmt.Errorf("ssh-chan error: %s", err)
	}
	go ssh.DiscardRequests(reqs)
	//remove on disconnect
	go u.unsetUDPChan(sshConn)
	//ready
	o := &udpChannel{
		r: gob.NewDecoder(rwc),
		w: gob.NewEncoder(rwc),
		c: rwc,
	}
	u.outbound = o
	u.Debugf("aquired channel")
	return o, nil
}

func (u *udpListener) unsetUDPChan(sshConn ssh.Conn) {
	sshConn.Wait()
	u.Debugf("lost channel")
	u.outboundMut.Lock()
	u.outbound = nil
	u.outboundMut.Unlock()
}

func (u *udpListener) monitorInactivity(ctx context.Context, cancel func()) error {
	previousRecv := atomic.LoadInt64(&u.recv)
	previousSent := atomic.LoadInt64(&u.sent)
	for {
		time.Sleep(INACTIVITY_CHECK_INTERVAL)
		shouldReturn := func() bool {
			u.remote.Lock()
			defer u.remote.Unlock()
			recv := atomic.LoadInt64(&u.recv)
			sent := atomic.LoadInt64(&u.sent)

			if time.Since(u.remote.LastTouched) > LAST_TOUCHED_TIMEOUT {
				if previousRecv == recv && previousSent == sent {
					return true
				}
				previousRecv = recv
				previousSent = sent
			}
			return false
		}()
		if shouldReturn {
			if settings.ClearFromActiveDynReverse(u.remote) {
				// We've cleared it from the cache, we'll continue monitoring the inactivity of the connection just in case it got sent just before the cache clear
				u.Infof("Cleared entry from active dynamic reverses")
				continue
			}

			u.Infof("Closing due to inactivity timeout")

			cancel()
			u.inbound.Close()
			if u.outbound != nil && u.outbound.c != nil {
				u.outbound.c.Close()
			}
			return nil
		}
	}
}
