package client

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/adamweixuan/getty"
	"net"
	"strings"
	"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))

	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)
		return
	}

	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)
			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))
	}

	NewTransfer(packet.SerialNumber, fmt.Sprintf("%s:%s", target, hostPort[1]), session, client, recorder).Start()
}

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.Data)
}

func (client *Client) handleDisconnectMessage(session getty.Session, packet *protocol.MajoraPacket) {
	log.Run().Debugf("[handleDisconnectMessage] %d->session closed %v extra:%s", packet.SerialNumber, session.IsClosed(), packet.Extra)
	// disconnect 会多回复一个关闭的包，暂时忽略，后续更改关闭交互流程
	t, ok := client.GetTransfer(packet.SerialNumber)
	if !ok {
		log.Run().Debug("handleDisconnectMessage can't find transfer")
		return
	}
	t.recorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("Received disconnect packet (sn:%d)", t.serialNumber))
	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 {
		return nil, false
	}
	transfer := t.(*Transfer)
	return transfer, true
}

func (client *Client) AddTransfer(sn int64, transfer *Transfer) error {
	if _, ok := client.transferStore.Load(sn); ok {
		return errors.New(fmt.Sprintf("Already has serialNumber:%d transfer", sn))
	}
	client.transferStore.Store(sn, transfer)
	return nil
}

func (client *Client) OnUpStreamConnectSuccess(t *Transfer) {
	t.recorder.RecordEvent(trace.ConnectEvent, fmt.Sprintf("Start replay natServer connect ready (sn:%d)", t.serialNumber))
	majoraPacket := protocol.TypeConnectReady.CreatePacket()
	majoraPacket.SerialNumber = t.serialNumber
	majoraPacket.Extra = client.config.ClientID
	if client.session.IsClosed() {
		log.Run().Warnf("[handleConnect] %d -> nat server is closed", t.serialNumber)
		t.recorder.RecordErrorEvent(trace.ConnectEvent, fmt.Sprintf("NatServer is closed (sn:%d)", t.serialNumber),
			nil)
		return
	}
	if _, _, err := client.session.WritePkg(majoraPacket, 0); err != nil {
		log.Error().Errorf("[handleConnect] %d->write pkg to nat server with error %+v", t.serialNumber,
			err)
		t.recorder.RecordErrorEvent(trace.ConnectEvent, fmt.Sprintf("Write pkg to natServer failed (sn:%d)",
			t.serialNumber), err)
		return
	}
	log.Run().Debugf("[handleConnect] %d->connect success to %s ", t.serialNumber, t.target)
	t.recorder.RecordEvent(trace.ConnectEvent, fmt.Sprintf("Replay natServer connect ready success (sn:%d)", t.serialNumber))
	err := client.AddTransfer(t.serialNumber, t)
	if err != nil {
		t.recorder.RecordErrorEvent(trace.UpStreamEvent, "Add transfer error: %+v", err)
		return
	}

}

func (client *Client) OnUpStreamConnectFailed(t *Transfer, err error) {
	t.recorder.RecordErrorEvent(trace.UpStreamEvent, "Connect upstream failed", err)
}

func (client *Client) OnUpStreamWriteError(t *Transfer, err error) {
	t.recorder.RecordErrorEvent(trace.DisconnectEvent, fmt.Sprintf("UpStream write error (sn:%d)", t.serialNumber), err)
	client.transferStore.Delete(t.serialNumber)
	t.Close()
	majoraPacket := protocol.TypeDisconnect.CreatePacket()
	majoraPacket.SerialNumber = t.serialNumber
	majoraPacket.Extra = client.config.ClientID
	if allCnt, sendCnt, err := client.session.WritePkg(majoraPacket, 0); err != nil {
		log.Run().Warnf("[closeVirtualConnection] ->%d error %s session closed %v allCnt %d sendCnt %d",
			t.serialNumber, err.Error(), client.session, allCnt, sendCnt)
		t.recorder.RecordErrorEvent(trace.DisconnectEvent,
			fmt.Sprintf("send disconnect to natServer failed closed:%v allCnt %d sendCnt %d (sn:%d)",
				client.session.IsClosed(), allCnt, sendCnt, t.serialNumber), err)
		return
	}
	t.recorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("Send disconnect to natServer success (sn:%d)", t.serialNumber))
}

func (client *Client) OnUpStreamReadError(t *Transfer, err error) {
	t.recorder.RecordErrorEvent(trace.DisconnectEvent, fmt.Sprintf("UpStream read error (sn:%d)", t.serialNumber), err)
	client.transferStore.Delete(t.serialNumber)
	t.Close()
	majoraPacket := protocol.TypeDisconnect.CreatePacket()
	majoraPacket.SerialNumber = t.serialNumber
	majoraPacket.Extra = client.config.ClientID
	if allCnt, sendCnt, err := client.session.WritePkg(majoraPacket, 0); err != nil {
		log.Run().Warnf("[closeVirtualConnection] ->%d error %s session closed %v allCnt %d sendCnt %d",
			t.serialNumber, err.Error(), client.session, allCnt, sendCnt)
		t.recorder.RecordErrorEvent(trace.DisconnectEvent,
			fmt.Sprintf("send disconnect to natServer failed closed:%v allCnt %d sendCnt %d (sn:%d)",
				client.session.IsClosed(), allCnt, sendCnt, t.serialNumber), err)
		return
	}
	t.recorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("Send disconnect to natServer success (sn:%d)", t.serialNumber))
}

func (client *Client) OnDownStreamWriteError(t *Transfer, err error) {
	t.recorder.RecordErrorEvent(trace.DisconnectEvent, fmt.Sprintf("DownStrwam  error (sn:%d)", t.serialNumber), err)
	client.Close()
	if !client.session.IsClosed() {
		client.session.Close()
	}
}
