package main

import (
	"io"
	"net"
	"os"
	"sync"

	"github.com/integrii/flaggy"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
)

func main() {
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})

	var (
		listenAddr  = ":8000"
		forwardAddr string
		mode        = "tcp"
	)

	flaggy.String(&listenAddr, "l", "listen", "Address to listen on")
	flaggy.String(&forwardAddr, "f", "forward", "Address to forward to")
	flaggy.String(&mode, "m", "mode", "Mode to use (tcp, udp)")
	flaggy.Parse()

	if forwardAddr == "" {
		log.Fatal().Msg("Forward address is required")
	}

	if mode == "tcp" {
		tcpListen(listenAddr, forwardAddr)
	} else if mode == "udp" {
		udpListen(listenAddr, forwardAddr)
	} else {
		log.Fatal().Msg("Invalid mode")
	}

}

func tcpListen(listenAddr, forwardAddr string) {
	var (
		ln  net.Listener
		err error
	)

	if ln, err = net.Listen("tcp", listenAddr); err != nil {
		log.Panic().Err(err).Str("Mode", "TCP").Msg("Failed to listen")
	}

	log.Info().Str("Mode", "TCP").Str("Listen Address", listenAddr).Msg("Listening")

	defer ln.Close()

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Err(err).Str("Mode", "TCP").Msg("Failed to accept connection")
			continue
		}
		log.Info().Str("Mode", "TCP").Str("Incoming Address", conn.RemoteAddr().String()).Msg("Connection accepted")

		go handleConnection(conn, forwardAddr)
	}
}

func udpListen(la, forwardAddr string) {
	// Listen for incoming UDP packets on port 3000
	listenAddr, err := net.ResolveUDPAddr("udp", la)
	if err != nil {
		log.Err(err).Str("Mode", "UDP").Msg("Failed to resolve UDP listen address")
		return
	}

	listener, err := net.ListenUDP("udp", listenAddr)
	if err != nil {
		log.Err(err).Str("Mode", "UDP").Msg("Failed to listen on UDP port")
		return
	}
	defer listener.Close()

	log.Info().Str("Mode", "UDP").Str("Listen Address", listenAddr.String()).Msg("Listening")

	// Forward incoming UDP packets to another host and port
	remoteAddr, err := net.ResolveUDPAddr("udp", forwardAddr)
	if err != nil {
		log.Err(err).Str("Mode", "UDP").Msg("Failed to resolve remote address")
		return
	}

	var (
		callerAddr *net.UDPAddr
		sendAddr   *net.UDPAddr
	)

	buf := make([]byte, 2048)

	for {

		n, addr, err := listener.ReadFromUDP(buf)
		if err != nil {
			log.Err(err).Str("Mode", "UDP").Msg("Failed to read from UDP connection")
			continue
		}

		if callerAddr == nil {
			callerAddr = addr
		}

		if addr.String() == remoteAddr.String() {
			sendAddr = callerAddr
		} else {
			sendAddr = remoteAddr
			callerAddr = addr
		}

		//log.Info().Str("Mode", "UDP").Str("Incoming Address", addr.String()).Str("Forward Address", sendAddr.String()).Msg("Packet received")

		_, err = listener.WriteToUDP(buf[:n], sendAddr)
		if err != nil {
			log.Err(err).Str("Mode", "udp").Msg("Failed to write to UDP connection")
			continue
		}

	}
}

func handleConnection(conn net.Conn, forwardAddr string) {
	defer conn.Close()

	fwrd, err := net.Dial("tcp", forwardAddr)
	if err != nil {
		log.Err(err).Str("Mode", "TCP").Msg("Failed to connect to forward address")
		return
	}

	var wg sync.WaitGroup

	wg.Add(2)

	go func() {
		defer fwrd.Close()
		defer wg.Done()

		_, err := io.Copy(fwrd, conn)
		if err != nil {
			log.Err(err).Str("Mode", "TCP").Msg("Failed to copy from connection to forward address")
		}
	}()

	go func() {
		defer conn.Close()
		defer wg.Done()

		_, err := io.Copy(conn, fwrd)
		if err != nil {
			log.Err(err).Str("Mode", "TCP").Msg("Failed to copy from forward address to connection")
		}
	}()

	wg.Wait()
	log.Info().Str("Mode", "TCP").Str("Incoming Address", conn.RemoteAddr().String()).Msg("Connection closed")
}