package app

import (
	"context"
	"git.makemake.in/kzkzzzz/mycommon/myconf"
	"git.makemake.in/kzkzzzz/mycommon/mylog"
	"io"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

type ProxyItem struct {
	RemoteAddr []string
	LocalAddr  string
}

var (
	mainWg              = &sync.WaitGroup{}
	mainCtx, mainCancel = context.WithCancel(context.Background())

	proxyInfo map[string]*ProxyItem
)

func Run() {
	err := myconf.Conf().UnmarshalKey("proxy", &proxyInfo)
	if err != nil {
		mylog.Error(err)
		return
	}

	names := make([]string, 0, len(proxyInfo))
	for name := range proxyInfo {
		names = append(names, name)
	}

	for _, name := range names {
		go listenTcp(name, proxyInfo[name])
	}

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

	v := <-ch
	mylog.Infof("捕获退出信号: %s", v.String())
	mainCancel()

	mainWg.Wait()
	mylog.Info("stop")

}

func listenTcp(name string, item *ProxyItem) {
	listen, err := net.Listen("tcp4", item.LocalAddr)
	if err != nil {
		mylog.Errorf("[%s] listen err: %s", name, err)
		return
	}

	mylog.Infof("[%s] forward %s -> %+v", name, item.LocalAddr, item.RemoteAddr)

	for {
		localConn, err := listen.Accept()
		if err != nil {
			mylog.Errorf("conn accept err: %s", item.LocalAddr, err)
			return
		}

		mainWg.Add(1)
		go proxyTcp(localConn, item)
	}
}

func proxyTcp(localConn net.Conn, item *ProxyItem) {
	defer mainWg.Done()

	randRemoteAddr := item.RemoteAddr[rand.Intn(len(item.RemoteAddr))]

	remoteConn, err := net.DialTimeout("tcp", randRemoteAddr, time.Second*3)
	if err != nil {
		localConn.Close()
		mylog.Errorf("connect remote err: %s", err)
		return
	}

	mylog.Debugf("start proxy %s -> %s", localConn.RemoteAddr(), randRemoteAddr)

	waitCh := make(chan struct{})
	go func() {
		wg0 := &sync.WaitGroup{}
		wg0.Add(2)
		go copyConn(wg0, localConn, remoteConn)
		go copyConn(wg0, remoteConn, localConn)
		wg0.Wait()
		close(waitCh)
	}()

	select {
	case <-waitCh:

	case <-mainCtx.Done():
		localConn.Close()
		remoteConn.Close()
	}

	mylog.Debugf("stop proxy %s -> %s", localConn.RemoteAddr(), randRemoteAddr)

}

func copyConn(wg0 *sync.WaitGroup, localConn, remoteConn net.Conn) {
	defer func() {
		wg0.Done()
		localConn.Close()
		remoteConn.Close()
	}()

	_, err := io.Copy(remoteConn, localConn)
	if err != nil {
		mylog.Debug(err)
	}

}