package client

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"github.com/adamweixuan/getty"
	"net"
	"strings"
	"sync"

	"virjar.com/majora-go/common"
	"virjar.com/majora-go/global"
	"virjar.com/majora-go/log"
	"virjar.com/majora-go/protocol"
	"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("Get map data from connect packet failed (sn:%d)", packet.SerialNumber)
	}
	sessionId, ok := m[trace.MajoraSessionName]
	if !ok {
		log.Error().Errorf("Get sessionId from connect packet failed (sn:%d)", packet.SerialNumber)
	}
	enableTrace, ok := m[trace.MajoraTraceEnable]
	if !ok {
		log.Error().Errorf("Get enableTrace from connect packet failed (sn:%d)", packet.SerialNumber)
	}
	user, ok := m[trace.MajoraSessionUser]
	if !ok {
		log.Error().Errorf("Get user from connect packet failed (sn:%d)", packet.SerialNumber)
	}
	recorder := trace.NewRecorder(sessionId, client.host, user, enableTrace == "true")
	recorder.RecordEvent(trace.ConnectEvent, fmt.Sprintf("Start handle connect to %s (sn:%d)",
		packet.Extra, packet.SerialNumber))

	if session.IsClosed() {
		log.Run().Warnf("[handleConnect] %d -> nat server is closed", packet.SerialNumber)
		recorder.RecordErrorEvent(trace.ConnectEvent, fmt.Sprintf("NatServer is closed (sn:%d)", packet.SerialNumber), nil)
		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)
		recorder.RecordErrorEvent(trace.ConnectEvent,
			fmt.Sprintf("Connect extra invalid %s (%d)", packet.Extra, packet.SerialNumber), nil)
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}

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

	target := hostPort[0]
	ip, err := client.dnsCache.Get([]byte(hostPort[0]))
	if err != nil {
		recorder.RecordEvent(trace.DnsResolveEvent, fmt.Sprintf("Dns cache miss %s ", hostPort[0]))
		hosts, dnsErr := net.LookupHost(hostPort[0])
		if dnsErr != nil {
			recorder.RecordErrorEvent(trace.DnsResolveEvent, fmt.Sprintf("Resolve %s ip error", hostPort[0]), dnsErr)
			client.closeVirtualConnection(session, packet.SerialNumber)
			return
		}
		err := client.dnsCache.Set([]byte(hostPort[0]), []byte(hosts[0]), int(global.Config.DnsCacheDuration.Seconds()))
		if err != nil {
			recorder.RecordErrorEvent(trace.DnsResolveEvent, fmt.Sprintf("Dns cache set error %s", hostPort[0]), err)
		}
		target = hosts[0]
	} else {
		target = string(ip)
		recorder.RecordEvent(trace.DnsResolveEvent, fmt.Sprintf("Dns cache hit %s -> %s", hostPort[0], target))
	}

	conn, err := dialer.Dial(common.TCP, fmt.Sprintf("%s:%s", target, hostPort[1]))
	if err != nil {
		log.Error().Errorf("[handleConnect] %d->connect to %s->%s", packet.SerialNumber, packet.Extra, err.Error())
		recorder.RecordErrorEvent(trace.ConnectEvent,
			fmt.Sprintf("Connect to %s failed (sn:%d)", packet.Extra, packet.SerialNumber), err)
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	}
	tcpConn := conn.(*net.TCPConn)
	_ = tcpConn.SetNoDelay(true)
	_ = tcpConn.SetKeepAlive(true)

	t := NewTransfer(packet.SerialNumber, client, tcpConn, recorder)
	t.SetTransferToUpstreamFunc(client.transferToUpstream)
	t.SetTransferToDownstreamFunc(client.transferToDownstream)
	client.AddTransfer(packet.SerialNumber, t, packet.Extra)
	recorder.RecordEvent(trace.ConnectEvent, fmt.Sprintf("Connect to %s success, local: %s -> remote:%s (aaasn:%d)",
		packet.Extra, tcpConn.LocalAddr(), tcpConn.RemoteAddr(), packet.SerialNumber))

	recorder.RecordEvent(trace.ConnectEvent, fmt.Sprintf("Start replay natServer connect ready (sn:%d)", packet.SerialNumber))
	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)
		recorder.RecordErrorEvent(trace.ConnectEvent, fmt.Sprintf("NatServer is closed (sn:%d)", packet.SerialNumber),
			nil)
		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())
		recorder.RecordErrorEvent(trace.ConnectEvent, fmt.Sprintf("Write pkg to natServer failed (sn:%d)",
			packet.SerialNumber), err)
		client.closeVirtualConnection(session, packet.SerialNumber)
		return
	} else {
		t.Start()
		log.Run().Debugf("[handleConnect] %d->connect success to %s ", packet.SerialNumber, packet.Extra)
		recorder.RecordEvent(trace.ConnectEvent, fmt.Sprintf("Replay natServer connect ready success (sn:%d)", packet.SerialNumber))
	}
}

func decodeMap(data []byte) map[string]string {
	result := make(map[string]string, 2)
	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) {
	t, ok := client.GetTransfer(packet.SerialNumber)
	if !ok {
		log.Error().Errorf("can't find transfer")
		return
	}
	t.recorder.RecordEvent(trace.TransferEvent,
		fmt.Sprintf("Receive transfer packet from natServer,start to be forward to target, len:%d (%d)", len(packet.Data), packet.SerialNumber))
	t.TransferToUpstream(packet)
}

func (client *Client) handleDisconnectMessage(session getty.Session, packet *protocol.MajoraPacket) {
	log.Run().Debugf("[handleDisconnectMessage] %d->session closed %v extra:%s", packet.SerialNumber, session.IsClosed())
	t, ok := client.GetTransfer(packet.SerialNumber)
	if !ok {
		log.Error().Errorf("handleDisconnectMessage can't find transfer")
		return
	}
	t.Close()

}

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

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

func (client *Client) GetRecorder(sn int64) trace.Recorder {
	t, ok := client.transferStore.Load(sn)
	if !ok {
		log.Run().Warnf("[GetRecorder] get session failed, maybe already closed (%d)", sn)
		return trace.NewRecorder("", "", "", true)
	}
	transfer := t.(*Transfer)
	return transfer.recorder

}

func (client *Client) GetTransfer(sn int64) (*Transfer, bool) {
	t, ok := client.transferStore.Load(sn)
	if !ok {
		log.Error().Errorf("[GetTransfer] error, not exist (sn:%d)", sn)
		return nil, false
	}
	transfer := t.(*Transfer)
	return transfer, true
}

func (client *Client) AddTransfer(sn int64, transfer *Transfer, addr string) {
	if _, ok := client.transferStore.Load(sn); ok {
		log.Error().Errorf("[AddTransfer] %d->error, has one", sn)
	}
	client.transferStore.Store(sn, transfer)
	log.Run().Debugf("[AddTransfer] %d->%s success", sn, 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)
		}
	}()
	client.closeVirtualConnection(natSession, serialNumber)
	_ = upStreamSession.Close()
}

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

	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("Start send disconnect to natServer (sn:%d)", 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("send disconnect to natServer failed closed:%v allCnt %d sendCnt %d (sn:%d)",
				session.IsClosed(), allCnt, sendCnt, serialNumber), err)
		session.Close()
	}
	client.transferStore.Delete(serialNumber)
	traceRecorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("Send disconnect to natServer success (sn:%d)", serialNumber))
}

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

func (client *Client) transferToUpstream(t *Transfer, p *protocol.MajoraPacket) {
	cnt, err := t.upstreamConn.Write(p.Data)
	if err != nil {
		log.Error().Errorf("[handleTransfer] %d->write to upstream fail for %s", p.SerialNumber, err)
		traceMessage := fmt.Sprintf("Write to upstream failed (%d)", p.SerialNumber)
		t.recorder.RecordErrorEvent(trace.TransferEvent, traceMessage, err)
		client.closeVirtualConnection(client.session, p.SerialNumber)
		return
	}

	if cnt != len(p.Data) {
		log.Error().Errorf("[handleTransfer] %d-> write not all data for expect->%d/%d",
			p.SerialNumber, len(p.Data), cnt)
		traceMessage := fmt.Sprintf("Write not all data for expect -> %d/%d (sn:%d)", len(p.Data), cnt, p.SerialNumber)
		t.recorder.RecordErrorEvent(trace.TransferEvent, traceMessage, nil)
		client.closeVirtualConnection(client.session, p.SerialNumber)
		return
	}
	log.Run().Debugf("[handleTransfer] %d-> success dataLen: %d", p.SerialNumber, len(p.Data))
	traceMessage := fmt.Sprintf("transfer data success (%d)", p.SerialNumber)
	t.recorder.RecordEvent(trace.TransferEvent, traceMessage)
}

func (client *Client) transferToDownstream(t *Transfer, data []byte, err error) {
	if err != nil {
		opErr, ok := err.(*net.OpError)
		if ok && opErr.Err.Error() == "i/o timeout" {
			recorderMessage := fmt.Sprintf("Upstream deadDeadline start close (sn:%d)", t.serialNumber)
			t.recorder.RecordEvent(trace.UpStreamEvent, recorderMessage)
		} else {
			log.Run().Debugf("[handleUpStream] %d->read with error:%+v,l:%s->r:%s",
				t.serialNumber, err, t.upstreamConn.LocalAddr(), t.upstreamConn.RemoteAddr())
			recorderMessage := fmt.Sprintf("Read with l:%s->r:%s (sn:%d) ",
				t.upstreamConn.LocalAddr(), t.upstreamConn.RemoteAddr(), t.serialNumber)
			t.recorder.RecordErrorEvent(trace.UpStreamEvent, recorderMessage, err)
		}
		t.client.OnClose(t.client.session, t.upstreamConn, t.serialNumber)
	} else {
		t.recorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("read count: %d (sn:%d)",
			len(data), t.serialNumber))

		t.recorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("Start write to natServer (sn:%d)", t.serialNumber))
		pack := protocol.TypeTransfer.CreatePacket()
		pack.Data = data
		pack.SerialNumber = t.serialNumber
		if _, _, err := t.client.session.WritePkg(pack, 0); err != nil {
			log.Error().Errorf("[handleUpStream] %d-> write to server fail %+v", t.client, err.Error())
			t.recorder.RecordErrorEvent(trace.UpStreamEvent,
				fmt.Sprintf("Write to natServer failed (sn:%d)", t.serialNumber), err)
			t.client.OnClose(t.client.session, t.upstreamConn, t.serialNumber)
		} else {
			log.Run().Debugf("[handleUpStream] %d->success dataLen:%d ", t.serialNumber, len(pack.Data))
			t.recorder.RecordEvent(trace.UpStreamEvent,
				fmt.Sprintf("Write to natServer success(sn:%d)", pack.SerialNumber))
		}
	}
}
