package main

import (
	"net"
	"./log"
	"time"

	"sync"

	"encoding/binary"

	"unsafe"

	"sync/atomic"

	"io"

	"./pb"

	"crypto/md5"

	"strconv"

	"fmt"

	"math"

	"github.com/golang/protobuf/proto"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

type Acceptor interface {
	Start(address string) error
	Close()
}

func NewAcceptor(log *log.Logger, onAccept func(conn net.Conn)) Acceptor {
	return &socketAcceptor{
		log:        log,
		onAccept:   onAccept,
		blockIPMap: map[string]int64{},
	}
}

type socketAcceptor struct {
	listener   net.Listener
	close      chan struct{}	// 执行退出
	closeDone  chan struct{}	// 用于同步退出完成
	log        *log.Logger		// 日志
	onAccept   func(conn net.Conn)
	isRunning  bool			// 是否在运行
	blockIPMap map[string]int64
	sync.Mutex
}

func (acceptor *socketAcceptor) Start(address string) error {
	acceptor.Lock()
	defer acceptor.Unlock()
	if acceptor.isRunning {
		return nil
	}
	ln, err := listen(address)	// 返回 listener
	if err != nil {
		acceptor.log.Error("listen failed", zap.String("addr", address), zap.Error(err))
		return err
	}
	acceptor.close = make(chan struct{})
	acceptor.closeDone = make(chan struct{})
	acceptor.isRunning = true
	acceptor.listener = ln
	go acceptor.accept()
	return nil
}

func (acceptor *socketAcceptor) Close() {
	select {
	case <-acceptor.close:
		return
	default:
		close(acceptor.close)
		acceptor.listener.Close()
		<-acceptor.closeDone
	}
}

// run in goroutine
func (acceptor *socketAcceptor) accept() {
	var tempDelay time.Duration
	for {
		conn, e := acceptor.listener.Accept()
		if e != nil {
			select {
			case <-acceptor.close:	//a 已停止工作
				goto stopAccept
			default:
			}
			//b 如果没有关闭，则睡眠后继续
			if ne, ok := e.(net.Error); ok && ne.Temporary() {
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond	// 1 第一次睡眠 5 ms
				} else {
					tempDelay *= 2 // 2 第二次以后翻倍
				}
				// 3 但最多休眠 1s
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				acceptor.log.Error("tcp: Accept error,will retrying", zap.Duration("retrying after", tempDelay),
					zap.String("addr", acceptor.listener.Addr().String()), zap.Error(e))
				time.Sleep(tempDelay)
				continue
			}
			//c 否则退出
			acceptor.log.Error("tcp: Accept error,will stop accept", zap.String("addr", acceptor.listener.Addr().String()), zap.Error(e))
			goto stopAccept
		}
		//tempDelay = 0
		//now := timestamp
		//a, ok := conn.RemoteAddr().(*net.TCPAddr)
		//if !ok {
		//	conn.Close()
		//	continue
		//}
		//ip := a.IP.String()
		//if v, ok := acceptor.blockIPMap[ip]; ok {
		//	if v == now {
		//		conn.Close()
		//		atomic.AddInt32(&acceptBlockCount, 1)
		//		continue
		//	}
		//}
		//acceptor.blockIPMap[ip] = now
		acceptor.onAccept(conn)	// 无错误执行本函数
	}// for

stopAccept:
	acceptor.log.Info("tcp: stopped accept", zap.String("addr", acceptor.listener.Addr().String()))
	close(acceptor.closeDone)
	acceptor.isRunning = false
}

type bufferPool struct {
	p *sync.Pool
}

var defaultBuffPool = bufferPool{
	p: &sync.Pool{
		New: func() interface{} {
			return make([]byte, 1024*2)	//2k缓冲
		},
}}

func (b *bufferPool) Get() []byte {
	return b.p.Get().([]byte)
}
func (b *bufferPool) put(data []byte) {	// 小写？
	b.p.Put(data)
}

type proxy struct {
	sync.Mutex
	isNewPacket bool
	listenAddr  string
	acceptor    Acceptor
	container   *IPContainer
	proxyInfo   *proxyInfo
	privateKey  unsafe.Pointer
	md5Map      unsafe.Pointer
	packetCount sync.Map
}

func (p *proxy) ChangePacketCount(packet string, add bool) int32 {
	if add {
		count := new(int32)
		*count = 1
		v, ok := p.packetCount.LoadOrStore(packet, count)
		if !ok {
			return atomic.LoadInt32(count)
		}
		return atomic.AddInt32(v.(*int32), 1)
	} else {
		v, ok := p.packetCount.Load(packet)
		if ok {
			return atomic.AddInt32(v.(*int32), -1)
		} else {
			log.Error("no packet count", zap.String("packet", packet))
			return 0
		}
	}
}

func newProxy(isNewPacket bool, addr string, privateKey string, container *IPContainer, proxyInfo *proxyInfo) *proxy {
	p := &proxy{
		isNewPacket: isNewPacket,
		listenAddr:  addr,
		container:   container,
		proxyInfo:   proxyInfo,
	}
	p.packetCount = sync.Map{}
	privateKeyPtr := unsafe.Pointer(&privateKey)
	atomic.StorePointer(&p.privateKey, privateKeyPtr)
	p.acceptor = NewAcceptor(log.GetDefaultLog(), p.doAccept)
	md5Map := &sync.Map{}
	p.md5Map = unsafe.Pointer(md5Map)
	go p.clearMd5Map()
	return p
}

func (p *proxy) Run() {
	p.acceptor.Start(p.listenAddr)
}

func (p *proxy) ReListen(addr string) {
	p.acceptor.Close()
	p.acceptor = NewAcceptor(log.GetDefaultLog(), p.doAccept)
	p.acceptor.Start(addr)
}

func (p *proxy) SetNewPacket(newPacket bool) {
	p.isNewPacket = newPacket
}

func (p *proxy) SetPrivateKey(privateKey string) {
	privateKeyPtr := unsafe.Pointer(&privateKey)
	atomic.StorePointer(&p.privateKey, privateKeyPtr)
}

func (p *proxy) GetPrivateKey() *string {
	return (*string)(atomic.LoadPointer(&p.privateKey))
}

func (p *proxy) doAccept(conn net.Conn) {
	defaultWorkPool.put(func(timer *time.Timer) {
		p.newPacket(conn)
	})
}

func (p *proxy) getProxyConn() (net.Conn, string) {
	for {
		addr, ok := p.proxyInfo.getOneProxyAddr()
		if !ok {
			log.Error("close connection ,due to no proxy addr can use")
			return nil, ""
		}
		proxyConn, err := net.DialTimeout("tcp", addr, time.Second*3)
		if err != nil {
			log.Error("connect downstream failed", zap.String("addr", addr), zap.Error(err))
			return nil, ""
		} else {
			return proxyConn, addr
		}
	}
}

var defaultWorkPool = newWorkPool(10000)
var Md5NotEqual int32
var validateFailCount int32
var acceptBlockCount int32
var onlineCount int32

func init() {
	go func() {
		for {
			time.Sleep(time.Second)
			fmt.Println(atomic.LoadInt32(&Md5NotEqual), atomic.LoadInt32(&validateFailCount), atomic.LoadInt32(&acceptBlockCount), atomic.LoadInt32(&onlineCount))
		}
	}()
	go WarnLog()
}

func (p *proxy) newPacket(conn net.Conn) {
	defer conn.Close()
	//IP, ok := conn.RemoteAddr().(*net.TCPAddr)
	//if !ok {
	//	return
	//}
	//ip := IP.IP.String()
	//if p.container.IPIncrease(ip) {
	//	atomic.AddInt32(&OKCount, 1)
	//	defer p.container.IPDecrease(ip)
	//} else {
	//	atomic.AddInt32(&CloseCount, 1)
	//	return
	//}

	ok, packet := p.checkConn(conn)
	if !ok {
		return
	}
	proxyProxy, addr := p.getProxyConn()
	if proxyProxy == nil {
		return
	}
	atomic.AddInt32(&onlineCount, 1)
	log.Info("connect downstream success", zap.String("addr", addr), zap.String("client", conn.RemoteAddr().String()),
		zap.Int32("online count", atomic.LoadInt32(&onlineCount)), zap.String("packet", packet), zap.Int32("packetCount", p.ChangePacketCount(packet, true)))
	p.doHandler(conn, proxyProxy, addr, packet)
	atomic.AddInt32(&onlineCount, -1)
}

func (p *proxy) checkConn(conn net.Conn) (bool, string) {
	done := make(chan struct{})
	defer func() {
		select {
		case <-done:
		default:
			close(done)
		}
	}()

	defaultWorkPool.put(func(timer *time.Timer) {
		timer.Reset(time.Second * 5)
		select {
		case <-timer.C:
			warnLog.Warn("conn close by timeout")
			conn.Close()
		case <-done:
		}
	})

	sendBuffer := defaultBuffPool.Get()
	defer defaultBuffPool.put(sendBuffer)
	n, err := io.ReadFull(conn, sendBuffer[:4])
	if err != nil {
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}
	if n != 4 {
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}
	lenData := binary.LittleEndian.Uint16(sendBuffer[0:2])
	id := binary.LittleEndian.Uint16(sendBuffer[2:4])
	if int(lenData) >= len(sendBuffer)-4 {
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}

	if pb.C_FirstPacketCheck_MsgID(id) != pb.C_FirstPacketCheck_ID {
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}

	n, err = io.ReadFull(conn, sendBuffer[4:lenData])
	if err != nil {
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}

	pkg := &pb.C_FirstPacketCheck{}
	err = proto.Unmarshal(sendBuffer[4:lenData], pkg)
	if err != nil {
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}
	packet := pkg.GetPackage()
	if packet == "" {
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}
	sum := pkg.GetSum()
	if sum == "" {
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}
	now := time.Now().Unix()
	clientNow := pkg.GetTimestamp()

	data := packet + strconv.Itoa(int(clientNow)) + "asd1[;sd2'fsa1sx._=912sa1sas321asdqwz"
	md5Data, err := MD5(data)
	if err != nil {
		PutWarnLog(func() {
			warnLog.Warn("conn close by MD5 error", zap.String("error", err.Error()))
		})
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}
	if md5Data != sum {
		PutWarnLog(func() {
			warnLog.Warn("conn close by MD5 not equal", zap.String("server md5", md5Data), zap.String("client md5", sum))
		})
		atomic.AddInt32(&validateFailCount, 1)
		atomic.AddInt32(&Md5NotEqual, 1)
		return false, ""
	}
	if math.Abs(float64(now-clientNow)) > 60*60*2 {
		PutWarnLog(func() {
			warnLog.Warn("conn close by Time abs > 2 hour", zap.Int64("now", now), zap.Int64("clientTime", clientNow), zap.String("addr", conn.RemoteAddr().String()))
		})
		atomic.AddInt32(&validateFailCount, 1)
		return false, ""
	}
	if math.Abs(float64(now-clientNow)) > 10 {
		PutWarnLog(func() {
			warnLog.Warn("conn client time abs >10", zap.Int64("now", now), zap.Int64("clientTime", clientNow), zap.String("addr", conn.RemoteAddr().String()))
		})
	}
	if _, ok := p.getMd5Map().LoadOrStore(sum, struct{}{}); ok {
		PutWarnLog(func() {
			warnLog.Warn("conn close by same md5", zap.String("md5", sum))
		})
		return false, ""
	}
	return true, packet
}

func (p *proxy) clearMd5Map() {
	for {
		time.Sleep(time.Hour)
		md5Map := &sync.Map{}
		md5MapPtr := unsafe.Pointer(md5Map)
		atomic.StorePointer(&p.md5Map, md5MapPtr)
		PutWarnLog(func() {
			warnLog.Info("clear MD5 map success")
		})
	}
}
func (p *proxy) getMd5Map() *sync.Map {
	return (*sync.Map)(atomic.LoadPointer(&p.md5Map))
}

var warnLog = log.NewLogger(true, "./head_packet_validate", zapcore.DebugLevel)
var WarnChan = make(chan func(), 10000)

func PutWarnLog(f func()) {
	if len(WarnChan) < 9000 {
		WarnChan <- f
	}
}

func WarnLog() {
	for f := range WarnChan {
		f()
	}
}

func MD5(src string) (string, error) {
	h := md5.New()
	_, err := io.WriteString(h, src)
	if err != nil {
		return "", err
	}
	cipherStr := h.Sum(nil)
	return fmt.Sprintf("%x", cipherStr), nil
}

type CloseReason struct {
	ip     string
	closer string
	err    string
	sync.Mutex
}

func (cr *CloseReason) String() string {
	cr.Lock()
	defer cr.Unlock()
	return cr.closer + "(" + cr.ip + ")(" + cr.err + ")"
}

func (cr *CloseReason) SetValue(ip, closer, err string) {
	cr.Lock()
	defer cr.Unlock()
	if cr.ip != "" {
		return
	}
	cr.ip = ip
	cr.closer = closer
	cr.err = err
}

func (p *proxy) doHandler(conn net.Conn, proxyConn net.Conn, proxyAddr string, packet string) {
	defer func() {
		if e := recover(); e != nil {
			log.Error("recover panic", zap.Any("error", e))
		}
	}()
	var err error
	closeCh := make(chan struct{})
	closeReason := &CloseReason{}
	go func(cr *CloseReason) {
		defer func() {
			if e := recover(); e != nil {
				log.Error("recover panic", zap.Any("error", e))
			}
		}()

		sendBuffer := defaultBuffPool.Get()
		for {
			nr, er := conn.Read(sendBuffer)
			if nr > 0 {
				nw, ew := proxyConn.Write(sendBuffer[0:nr])
				if ew != nil {
					closeReason.SetValue(proxyConn.RemoteAddr().String(), "CloseServerWrite", ew.Error())
					err = ew
					break
				}
				if nr != nw {
					closeReason.SetValue(proxyConn.RemoteAddr().String(), "CloseServerWrite", "nr!=nw")
					break
				}
			}
			if er != nil {
				closeReason.SetValue(conn.RemoteAddr().String(), "CloseClientRead", er.Error())
				break
			}
		}
		defaultBuffPool.put(sendBuffer)
		select {
		case <-closeCh:
		default:
			close(closeCh)
			proxyConn.Close()
		}
	}(closeReason)

	now := time.Now()
	writeBuffer := defaultBuffPool.Get()
	for {
		nr, er := proxyConn.Read(writeBuffer)
		if nr > 0 {
			nw, ew := conn.Write(writeBuffer[0:nr])
			if ew != nil {
				closeReason.SetValue(conn.RemoteAddr().String(), "CloseClientWrite", ew.Error())
				err = ew
				break
			}
			if nr != nw {
				closeReason.SetValue(conn.RemoteAddr().String(), "CloseClientWrite", "nr != nw ")
				break
			}
		}
		if er != nil {
			closeReason.SetValue(proxyConn.RemoteAddr().String(), "CloseServerRead", er.Error())
			break
		}
	}
	defaultBuffPool.put(writeBuffer)
	select {
	case <-closeCh:
	default:
		close(closeCh)
		proxyConn.Close()
	}
	log.Info("connect closed", zap.String("srcConnAddr", conn.RemoteAddr().String()), zap.String("desConnAddr", proxyAddr),
		zap.Duration("online seconds", time.Since(now)), zap.Int32("onlineCount", atomic.LoadInt32(&onlineCount)-1),
		zap.String("CloseReason", closeReason.String()), zap.String("packet", packet), zap.Int32("packetCount", p.ChangePacketCount(packet, false)))
}

func ListenTest(addr string) bool {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		log.Error("test listen failed", zap.String("addr", addr), zap.Error(err))
		return false
	}
	ln.Close()
	return true
}
