package breacher

import (
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/spf13/cobra"
)

var forwardCmd = &cobra.Command{
	Use:   "forward [from address] [to address]",
	Short: "port forwarding for UDP->UDP / TCP->TCP",
	Long: `port forwarding, defaults to use tcp if no flags given
ex.
breacher forward :8080 kumoly.io:5080
breacher forward :8080 :8000
breacher forward --udp :8080 192.168.51.211:53
`,
	Args: cobra.ExactArgs(2),
	Run: func(cmd *cobra.Command, args []string) {
		if ftcp {
			ForwardTCP(args[0], args[1])
		}
		if fudp {
			ForwardUDP(args[0], args[1])
		}
		if !(ftcp && fudp) {
			ForwardTCP(args[0], args[1])
		}
		stop := make(chan os.Signal, 1)
		signal.Notify(stop, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
		<-stop
	},
}

var (
	ftcp bool
	fudp bool
)

func init() {
	forwardCmd.Flags().BoolVar(&ftcp, "tcp", false, "use tcp protocol (default if no flags given)")
	forwardCmd.Flags().BoolVar(&fudp, "udp", false, "use udp protocol")
}

func ForwardTCP(from, to string) {
	ln, err := net.Listen("tcp", from)
	if err != nil {
		log.Println(err)
		os.Exit(1)
	}

	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				log.Println(err)
				continue
			}

			go func(con net.Conn) {
				proxy, err := net.Dial("tcp", to)
				if err != nil {
					log.Println(err)
					return
				}

				log.Println("forwarding tcp from", con.RemoteAddr())
				go exchange(conn, proxy)
				go exchange(proxy, conn)
			}(conn)
		}
	}()
}

func exchange(src, dest net.Conn) {
	defer src.Close()
	defer dest.Close()
	io.Copy(src, dest)
}

func ForwardUDP(from, to string) {
	connIn, err := net.ListenPacket("udp", from)
	if err != nil {
		log.Println(err)
		return
	}
	go func() {
		type pipeCache struct {
			Pipe  *io.PipeWriter
			Ready *uintptr
			TTL   time.Time
		}
		type hashableAddr struct {
			Network string
			String  string
		}
		pipes := make(map[hashableAddr]pipeCache)
		pipesLock := new(sync.RWMutex)
		go func() {
			for {
				time.Sleep(59 * time.Second)
				now := time.Now()
				for k, v := range pipes {
					if v.TTL.Before(now) {
						pipesLock.Lock()
						delete(pipes, k)
						pipesLock.Unlock()
						v.Pipe.Close()
					}
				}
			}
		}()
		buffer := make([]byte, 65537)
		for {
			packetLen, addrIn, err := connIn.ReadFrom(buffer)
			if err != nil {
				log.Println(err)
				if errNet, ok := err.(net.Error); ok {
					if errNet.Temporary() {
						log.Println(err)
						continue
					}
				}
				log.Fatalln(err)
			}
			log.Println("forwarding udp from", addrIn)
			pipesLock.RLock()
			if pipeOut, ok := pipes[hashableAddr{
				Network: addrIn.Network(),
				String:  addrIn.String(),
			}]; ok {
				pipesLock.RUnlock()
				pipeOut.TTL = time.Now().Add(180 * time.Second)
				if atomic.LoadUintptr(pipeOut.Ready) != 0 {
					pipeOut.Pipe.Write(buffer[:packetLen])
				}
			} else {
				pipesLock.RUnlock()
				firstPacket := make([]byte, packetLen)
				copy(firstPacket, buffer)
				go func(addrIn net.Addr, firstPacket []byte) {
					connOut, err := net.Dial("udp", to)
					var connWait sync.WaitGroup
					connWait.Add(2)
					if err != nil {
						log.Println(err)
						return
					}
					pipeIn, pipeOut := io.Pipe()
					ready := new(uintptr)
					pipe := pipeCache{
						Pipe:  pipeOut,
						Ready: ready,
						TTL:   time.Now().Add(180 * time.Second),
					}
					pipesLock.Lock()
					pipes[hashableAddr{
						Network: addrIn.Network(),
						String:  addrIn.String(),
					}] = pipe
					pipesLock.Unlock()
					go func() {
						var err error
						var packetLen int
						buffer := make([]byte, 65537)
						for {
							atomic.StoreUintptr(ready, 1)
							packetLen, err = pipeIn.Read(buffer)
							atomic.StoreUintptr(ready, 0)
							if err != nil {
								break
							}
							_, err = connOut.Write(buffer[:packetLen])
							if err != nil {
								break
							}
						}

						if err != io.EOF {
							log.Println(err)
						}
						pipesLock.Lock()
						delete(pipes, hashableAddr{
							Network: addrIn.Network(),
							String:  addrIn.String(),
						})
						pipesLock.Unlock()
						pipeIn.Close()
						if connOutTCP, ok := connOut.(*net.TCPConn); ok {
							connOutTCP.CloseWrite()
						} else {
							connOut.Close()
						}
						connWait.Done()
					}()
					go func() {
						var err error
						var packetLen int
						buffer := make([]byte, 65537)
						for {
							connOut.SetReadDeadline(time.Now().Add(180 * time.Second))
							packetLen, err = connOut.Read(buffer)
							if err != nil {
								break
							}
							_, err = connIn.WriteTo(buffer[:packetLen], addrIn)
							if err != nil {
								break
							}
						}
						if err != io.EOF {
							log.Println(err)
						}
						if connOutTCP, ok := connOut.(*net.TCPConn); ok {
							connOutTCP.CloseRead()
						}
						connWait.Done()
					}()
					pipeOut.Write(firstPacket)
					connWait.Wait()
					if connOutTCP, ok := connOut.(*net.TCPConn); ok {
						connOutTCP.Close()
					}
				}(addrIn, firstPacket)
			}
		}
	}()
}