package client

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"strings"
	"sync"

	"github.com/adamweixuan/getty"

	"virjar.com/majora-go/common"
	"virjar.com/majora-go/log"
	"virjar.com/majora-go/protocol"
	"virjar.com/majora-go/safe"
	"virjar.com/majora-go/trace"
)

var (
	HeartbeatPacket = protocol.TypeHeartbeat.CreatePacket()
	OfflinePacket   = protocol.TypeOffline.CreatePacket()
)

func (client *Client) handleHeartbeat(session getty.Session) {
	if _, _, err := session.WritePkg(HeartbeatPacket, common.HeartBeatTimeout); err != nil {
		log.Error().Errorf("handleHeartbeat error %+v %v", err, session.IsClosed())
	} else {
		log.Run().Infof("handleHeartbeat success")
	}
}

func (client *Client) handleConnect(packet *protocol.MajoraPacket, session getty.Session) {
	m := decodeMap(packet.Data)
	if len(m) <= 0 {
		log.Error().Errorf("SN: %d -> get connect map failed", packet.SerialNumber)
	}
	sessionId, ok := m[trace.MajoraSessionName]
	if !ok {
		log.Error().Errorf("SN: %d -> get sessionId failed", packet.SerialNumber)
	}
	traceSession := trace.NewSession(sessionId)
	session.SetAttribute(trace.MajoraSessionName, traceSession)
	traceSession.Recorder.RecordEvent(trace.ConnectEvent, "start handle connect")

	if session.IsClosed() {
		log.Run().Warnf("[handleConnect] %d -> nat server is closed", packet.SerialNumber)
		traceSession.Recorder.RecordEvent("[handleConnect]", fmt.Sprintf("SerialNumber: %d -> nat server is closed", packet.SerialNumber))
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}
	hostPort := strings.Split(packet.Extra, ":")
	if len(packet.Extra) == 0 || len(hostPort) != 2 {
		log.Error().Errorf("[handleConnect] invalid extra %s", packet.Extra)
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}

	dialer := net.Dialer{
		Timeout:   common.UpstreamTimeout,
		LocalAddr: client.localAddr,
	}

	conn, err := dialer.Dial(common.TCP, packet.Extra)
	if err != nil {
		log.Error().Errorf("[handleConnect] %d->connect to %s->%s", packet.SerialNumber, packet.Extra, err.Error())
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}

	tcpConn := conn.(*net.TCPConn)
	_ = tcpConn.SetNoDelay(true)
	_ = tcpConn.SetKeepAlive(true)
	client.AddConnection(packet, tcpConn, packet.Extra)
	majoraPacket := protocol.TypeConnectReady.CreatePacket()
	majoraPacket.SerialNumber = packet.SerialNumber
	majoraPacket.Extra = client.config.ClientID
	if session.IsClosed() {
		log.Run().Warnf("[handleConnect] %d -> nat server is closed", packet.SerialNumber)
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}
	if _, _, err := session.WritePkg(majoraPacket, 0); err != nil {
		log.Error().Errorf("[handleConnect] %d->write pkg to nat server with error %s", packet.SerialNumber,
			err.Error())
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	} else {
		log.Run().Debugf("[handleConnect] %d->connect success to %s ", packet.SerialNumber, packet.Extra)
		traceSession.Recorder.RecordEvent(trace.ConnectEvent, fmt.Sprintf("connect success to %s", packet.Extra))
	}
	client.handleUpStream(tcpConn, packet, session)
}

func decodeMap(data []byte) map[string]string {
	result := make(map[string]string, 1)
	var headerSize int8
	err := binary.Read(bytes.NewBuffer(data[:1]), binary.BigEndian, &headerSize)
	data = data[1:]
	if err != nil {
		return result
	}
	var i int8
	for i = 0; i < headerSize; i++ {
		var keyLength int8
		err = binary.Read(bytes.NewBuffer(data[:1]), binary.BigEndian, &keyLength)
		data = data[1:]

		key := make([]byte, keyLength)
		err = binary.Read(bytes.NewBuffer(data[:keyLength]), binary.BigEndian, &key)
		data = data[keyLength:]

		var valueLength int8
		err = binary.Read(bytes.NewBuffer(data[:1]), binary.BigEndian, &valueLength)
		data = data[1:]

		value := make([]byte, valueLength)
		err = binary.Read(bytes.NewBuffer(data[:valueLength]), binary.BigEndian, &value)
		data = data[valueLength:]

		result[string(key)] = string(value)
	}
	return result
}

func (client *Client) handleTransfer(packet *protocol.MajoraPacket, session getty.Session) {
	traceRecorder := trace.GetTraceRecorderFromSession(session)
	traceRecorder.RecordEvent(trace.TransferEvent, "receive transfer packet from natServer,need to be forward to target")

	load, ok := client.connStore.Load(packet.SerialNumber)
	if !ok {
		log.Error().Errorf("[handleTransfer] %d-> can not find connection", packet.SerialNumber)
		traceMessage := fmt.Sprintf("SerialNumber:%d -> can not find connection", packet.SerialNumber)
		traceRecorder.RecordErrorEvent(trace.TransferEvent, traceMessage, nil)
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}
	conn := load.(*net.TCPConn)
	cnt, err := conn.Write(packet.Data)

	if err != nil {
		log.Error().Errorf("[handleTransfer] %d->write to upstream fail for %s", packet.SerialNumber, err)
		traceMessage := fmt.Sprintf("SerialNumber:%d -> write to upstream fail", packet.SerialNumber)
		traceRecorder.RecordErrorEvent(trace.TransferEvent, traceMessage, nil)
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}

	if cnt != len(packet.Data) {
		log.Error().Errorf("[handleTransfer] %d-> write not all data for expect->%d/%d",
			packet.SerialNumber, len(packet.Data), cnt)
		traceMessage := fmt.Sprintf("SerialNumber: %d -> write not all data for expect -> %d/%d", packet.SerialNumber, len(packet.Data), cnt)
		traceRecorder.RecordErrorEvent(trace.TransferEvent, traceMessage, nil)
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}
	log.Error().Debugf("[handleTransfer] %d-> success %+v", packet.SerialNumber, string(packet.Data))
	traceMessage := fmt.Sprintf("SerialNumber: %d -> complete transfer", packet.SerialNumber)
	traceRecorder.RecordEvent(trace.TransferEvent, traceMessage)
}

func (client *Client) handleUpStream(conn *net.TCPConn, packet *protocol.MajoraPacket, session getty.Session) {
	safe.SageGo(func() {
		traceRecorder := trace.GetTraceRecorderFromSession(session)
		traceRecorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("SerialNumber:%d -> handleUpStream start", packet.SerialNumber))
		log.Run().Debugf("[handleUpStream] %d-> handleUpStream start...", packet.SerialNumber)
		for {
			buf := make([]byte, common.BufSize) // 4k
			cnt, err := conn.Read(buf)
			if err != nil {
				log.Run().Debugf("[handleUpStream] %d->read with error:%+v,l:%s->r:%s",
					packet.SerialNumber, err, conn.LocalAddr(), conn.RemoteAddr())
				recorderMessage := fmt.Sprintf("SerialNumber: %d -> read with l:%s->r:%s",
					packet.SerialNumber, conn.LocalAddr(), conn.RemoteAddr())
				traceRecorder.RecordErrorEvent(trace.UpStreamEvent, recorderMessage, err)
				client.OnClose(session, conn, packet.SerialNumber)
				break
			}
			traceRecorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("SerialNumber: %d read count: %d",
				packet.SerialNumber, cnt))
			pack := protocol.TypeTransfer.CreatePacket()
			pack.Data = buf[:cnt]
			pack.SerialNumber = packet.SerialNumber

			if _, _, err := session.WritePkg(pack, 0); err != nil {
				log.Error().Errorf("[handleUpStream] %d-> write to server fail %+v", packet.SerialNumber, err.Error())
				traceRecorder.RecordErrorEvent(trace.UpStreamEvent,
					fmt.Sprintf("SerialNumber: %d -> write to natServer failed", packet.SerialNumber), err)
				client.OnClose(session, conn, packet.SerialNumber)
				break
			} else {
				log.Run().Debugf("[handleUpStream] %d->success %+v", packet.SerialNumber, string(packet.Data))
				traceRecorder.RecordEvent(trace.UpStreamEvent,
					fmt.Sprintf("SerialNumber: %d write to natServer success", packet.SerialNumber))
			}
		}
	})
}

func (client *Client) handleDisconnectMessage(session getty.Session, packet *protocol.MajoraPacket) {
	traceRecorder := trace.GetTraceRecorderFromSession(session)
	traceRecorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("SN: %d start close session", packet.SerialNumber))
	log.Run().Debugf("[handleDisconnectMessage] %d->session closesd %v", packet.SerialNumber, session.IsClosed())
	if conn, ok := client.connStore.Load(packet.SerialNumber); ok {
		client.OnClose(session, conn.(net.Conn), packet.SerialNumber)
	}
}

func (client *Client) handleControlMessage(_ *protocol.MajoraPacket) {
	log.Run().Debugf("handleControlMessage")
}

// handleDestroyMessage 是直接关闭nat server ?
func (client *Client) handleDestroyMessage() {
	client.natTunnel.Close()
}

func (client *Client) AddConnection(packet *protocol.MajoraPacket, conn *net.TCPConn, addr string) {
	if _, ok := client.connStore.Load(packet.SerialNumber); ok {
		log.Error().Errorf("[AddConnection] %d->error, has one", packet.SerialNumber)
	}
	client.connStore.Store(packet.SerialNumber, conn)
	log.Run().Debugf("[AddConnection] %d->%s success", packet.SerialNumber, addr)
}

// OnClose 1. 本地缓存删除 2. 关闭连接  3. 通知natserver
func (client *Client) OnClose(natSession getty.Session, upStreamSession net.Conn, serialNumber int64) {
	defer func() {
		if err := recover(); err != nil {
			log.Error().Errorf("OnClose %+v", err)
		}
	}()
	traceRecorder := trace.GetTraceRecorderFromSession(natSession)
	client.connStore.Delete(serialNumber)
	_ = upStreamSession.Close()
	client.closeVirtualConnection(natSession, serialNumber)
	traceRecorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("SN: %d close success", serialNumber))
}

//closeVirtualConnection disconnect to server
func (client *Client) closeVirtualConnection(session getty.Session, serialNumber int64) {
	traceRecorder := trace.GetTraceRecorderFromSession(session)

	log.Run().Debugf("[closeVirtualConnection] %d->session closed %v", serialNumber, session.IsClosed())
	if session.IsClosed() {
		log.Run().Warnf("[closeVirtualConnection] %d->session is closed", serialNumber)
		return
	}

	traceRecorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("SN: %d -> start send disconnect to natServer", serialNumber))
	majoraPacket := protocol.TypeDisconnect.CreatePacket()
	majoraPacket.SerialNumber = serialNumber
	majoraPacket.Extra = client.config.ClientID
	if allCnt, sendCnt, err := session.WritePkg(majoraPacket, 0); err != nil {
		log.Run().Warnf("[closeVirtualConnection] ->%d error %s session closed %v allCnt %d sendCnt %d",
			serialNumber, err.Error(), session.IsClosed(), allCnt, sendCnt)
		traceRecorder.RecordErrorEvent(trace.DisconnectEvent,
			fmt.Sprintf("SN: %d -> send disconnect failed closed:%v allCnt %d sendCnt %d",
				serialNumber, session.IsClosed(), allCnt, sendCnt), err)
	}
	traceRecorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("SN: %d -> send disconnect to natServer success", serialNumber))
}

func (client *Client) CloseAll(session getty.Session) {
	defer func() {
		if err := recover(); err != nil {
			log.Error().Errorf("OnClose %+v", err)
		}
	}()
	client.connStore.Range(func(key, value interface{}) bool {
		serialNumber := key.(int64)
		conn, _ := value.(*net.TCPConn)
		log.Run().Debugf("[CloseAll] close serialNumber -> %d", serialNumber)
		client.OnClose(session, conn, serialNumber)
		return true
	})
	client.connStore = sync.Map{}
}
