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
	)

	flaggy.String(&listenAddr, "l", "listen", "Address to listen on")
	flaggy.String(&forwardAddr, "f", "forward", "Address to forward to")
	flaggy.Parse()

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

	ln, err := net.Listen("tcp", listenAddr)
	if err != nil {
		log.Panic().Err(err).Msg("Failed to listen")
	}
	defer ln.Close()

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

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

		go handleConnection(conn, forwardAddr)
	}

}

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

	fwrd, err := net.Dial("tcp", forwardAddr)
	if err != nil {
		log.Err(err).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).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).Msg("Failed to copy from forward address to connection")
		}
	}()

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