package net

import (
	"b612.me/starlog"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"strings"
	"time"
)

type NatTesterClient struct {
	MainPort  string `json:"mainport"`
	AltPort   string `json:"altport"`
	MainIP    string `json:"mainip"`
	AltIP     string `json:"altip"`
	RetryTime int    `json:"retrytime"`
	Timeout   int    `json:"timeout"`
	dns       []string
	ch        chan Message
}

type Message struct {
	Success bool
	Cmd     string
	Msg     string
	Address string
}

type NatType struct {
	Code           string
	RFC3489        string
	NintendoSwitch string
	Desc           string
}

func (n *NatTesterClient) GetMsg() Message {
	select {
	case msg := <-n.ch:
		return msg
	case <-time.After(time.Second * time.Duration(n.Timeout)):
		return Message{Success: false, Msg: "timeout"}
	}
}

func (n *NatTesterClient) RecvMsg(c *net.UDPConn) {
	for {
		buf := make([]byte, 1024)
		num, r, e := c.ReadFromUDP(buf)
		if e != nil {
			return
		}
		go n.Analyse(r, strings.Split(string(buf[:num]), "::"))
	}
}

func (n *NatTesterClient) Analyse(r *net.UDPAddr, cmds []string) {
	switch cmds[0] {
	case "ip":
		if len(cmds) == 2 {
			n.ch <- Message{Success: true, Cmd: "ip", Msg: cmds[1], Address: r.String()}
		}
	case "stage1":
		n.ch <- Message{Success: true, Cmd: "stage1", Msg: "stage1", Address: r.String()}
	case "stage2":
		n.ch <- Message{Success: true, Cmd: "stage2", Msg: "stage2", Address: r.String()}
	case "stage3":
		n.ch <- Message{Success: true, Cmd: "stage3", Msg: "stage3", Address: r.String()}
	}
}

func (n *NatTesterClient) Run() (NatType, error) {
	var firstAddr, secondAddr string
	var natTypeMap = map[string]NatType{
		"Block":     {Code: "Block", RFC3489: "Block", NintendoSwitch: "NAT F", Desc: "您的网络似乎禁止了UDP，无法连接到外部网络"},
		"Open":      {Code: "Open", RFC3489: "Open", NintendoSwitch: "Open", Desc: "您的网络为开放网络，拥有最好的上网体验"},
		"OpenBlock": {Code: "OpenBlock", RFC3489: "Symmetric Firewall", NintendoSwitch: "NAT F", Desc: "您的网络虽然时开放网络，但是存在对称型防火墙，可能会遇到严重的连接问题"},
		"NAT1":      {Code: "NAT1", RFC3489: "Full Cone", NintendoSwitch: "NAT A", Desc: "您的NAT类型为全锥形NAT，拥有最好的NAT体验"},
		"NAT2":      {Code: "NAT2", RFC3489: "Address Restricted Cone", NintendoSwitch: "NAT B", Desc: "您的NAT类型为地址限制锥形NAT，拥有良好的NAT体验"},
		"NAT3":      {Code: "NAT3", RFC3489: "Port Restricted Cone", NintendoSwitch: "NAT B/C", Desc: "您的NAT类型为端口限制锥形NAT，可能会遇到一些连接问题"},
		"NAT4":      {Code: "NAT4", RFC3489: "Symmetric", NintendoSwitch: "NAT C/D", Desc: "您的NAT类型为对称NAT，可能会遇到严重的连接问题"},
		"Unknown":   {Code: "Unknown", RFC3489: "Unknown", NintendoSwitch: "Unknown", Desc: "无法确定您的NAT类型"},
	}
	tmp, err := net.Dial("udp", n.MainIP+":80")
	if err != nil {
		return NatType{}, err
	}
	curIp := tmp.LocalAddr().(*net.UDPAddr).IP.String()
	starlog.Infof("Current Output IP: %s\n", curIp)
	localAddr, err := net.ResolveUDPAddr("udp", curIp+":0")
	if err != nil {
		return NatType{}, err
	}
	conn, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		return NatType{}, err
	}
	starlog.Infof("Listening on %s\n", conn.LocalAddr().String())
	defer conn.Close()
	go n.RecvMsg(conn)
	n.ch = make(chan Message)
	defer close(n.ch)
	succ := false
	mainAddr, err := net.ResolveUDPAddr("udp", n.MainIP+":"+n.MainPort)
	if err != nil {
		return NatType{}, err
	}
	altAddr, err := net.ResolveUDPAddr("udp", n.AltIP+":"+n.AltPort)
	if err != nil {
		return NatType{}, err
	}
	starlog.Noticef("Getting IP from NatServer Main\n")
	for i := 0; i < n.RetryTime; i++ {
		_, err = conn.WriteToUDP([]byte("ip"), mainAddr)
		if err != nil {
			starlog.Errorln("failed to get main ip,retrying:" + err.Error())
			continue
		}
		msg := n.GetMsg()
		if msg.Success && msg.Cmd == "ip" {
			starlog.Noticef("Remote IP: %s\n", msg.Address)
			starlog.Infof("Current IP: %s\n", msg.Msg)
			succ = true
			firstAddr = msg.Msg
			break
		}
		starlog.Errorln("failed to get main ip,retrying:" + msg.Msg)
	}
	if !succ {
		return NatType{}, fmt.Errorf("failed to get current ip")
	}
	{
		starlog.Noticef("Start NAT1 Test\n")
		succ = false
		for i := 0; i < n.RetryTime; i++ {
			_, err = conn.WriteToUDP([]byte("startnat1"), mainAddr)
			if err != nil {
				starlog.Errorln("failed to send nat1 test data,retrying:" + err.Error())
				continue
			}
			msg := n.GetMsg()
			if msg.Success && msg.Cmd == "stage1" {
				starlog.Noticef("Recv Nat1 Data From Remote IP: %s\n", msg.Address)
				succ = true
				break
			}
			starlog.Errorln("failed to recv Nat1 data,retrying:" + msg.Msg)
		}
		if succ {
			if strings.Split(firstAddr, ":")[0] == curIp {
				starlog.Infof("Current NAT Type: Open\n")
				return natTypeMap["Open"], nil
			}
			starlog.Infof("Current NAT Type: NAT1\n")
			return natTypeMap["NAT1"], nil
		} else {
			if strings.Split(firstAddr, ":")[0] == curIp {
				starlog.Infof("Current NAT Type: OpenBlock\n")
				return natTypeMap["OpenBlock"], nil
			}
		}
	}
	{
		starlog.Noticef("Start NAT2 Test\n")
		succ = false
		for i := 0; i < n.RetryTime; i++ {
			_, err = conn.WriteToUDP([]byte("startnat2"), mainAddr)
			if err != nil {
				starlog.Errorln("failed to send nat2 test data,retrying:" + err.Error())
				continue
			}
			msg := n.GetMsg()
			if msg.Success && msg.Cmd == "stage2" {
				starlog.Noticef("Recv Nat2 Data From Remote IP: %s\n", msg.Address)
				succ = true
				break
			}
			starlog.Errorln("failed to recv Nat2 data,retrying:" + msg.Msg)
		}
		if succ {
			starlog.Infof("Current NAT Type: NAT2\n")
			return natTypeMap["NAT2"], nil
		}
	}
	{
		starlog.Noticef("Start NAT3 Test\n")
		succ = false
		for i := 0; i < n.RetryTime; i++ {
			_, err = conn.WriteToUDP([]byte("startnat3"), mainAddr)
			if err != nil {
				starlog.Errorln("failed to send nat1 test data,retrying:" + err.Error())
				continue
			}
			msg := n.GetMsg()
			if msg.Success && msg.Cmd == "stage3" {
				starlog.Noticef("Recv Nat1 Data From Remote IP: %s\n", msg.Address)
				succ = true
				break
			}
			starlog.Errorln("failed to recv Nat3 data,retrying:" + msg.Msg)
		}
		if !succ {
			starlog.Errorf("Failed to get NAT Type\n")
			return natTypeMap["Unknown"], fmt.Errorf("failed to get nat type")
		}
	}
	succ = false
	starlog.Noticef("Gettting IP from NatServer Alt\n")
	for i := 0; i < n.RetryTime; i++ {
		_, err = conn.WriteToUDP([]byte("ip"), altAddr)
		if err != nil {
			starlog.Errorln("failed to get alt ip,retrying:" + err.Error())
			continue
		}
		msg := n.GetMsg()
		if msg.Success && msg.Cmd == "ip" {
			starlog.Noticef("Remote IP: %s\n", msg.Address)
			starlog.Infof("Current IP: %s\n", msg.Msg)
			succ = true
			secondAddr = msg.Msg
			break
		}
		starlog.Errorln("failed to get alt ip,retrying:" + msg.Msg)
	}
	if !succ {
		starlog.Errorf("Failed to get NAT Type\n")
		return natTypeMap["Unknown"], fmt.Errorf("failed to get nat type")
	}
	starlog.Debugf("First IP: %s, Second IP: %s\n", firstAddr, secondAddr)
	starlog.Debugf("Listening on %s\n", conn.LocalAddr().String())
	if firstAddr == secondAddr {
		starlog.Infof("Current NAT Type: NAT3\n")
		return natTypeMap["NAT3"], nil
	}
	starlog.Infof("Current NAT Type: NAT4\n")
	return natTypeMap["NAT4"], nil
}

func (n *NatTesterClient) ServeAndRun(addr string) (NatType, error) {
	starlog.SetShowFlag(false)
	starlog.SetShowFuncName(false)
	starlog.SetShowOriginFile(false)
	data, err := net.LookupTXT(addr)
	if err != nil {
		return NatType{}, err
	}
	if len(data) == 0 {
		return NatType{}, fmt.Errorf("no data found")
	}
	err = json.Unmarshal([]byte(data[0]), n)
	if err != nil {
		return NatType{}, err
	}
	starlog.Debugf("MainIP: %s, MainPort: %s, AltIP: %s, AltPort: %s\n", n.MainIP, n.MainPort, n.AltIP, n.AltPort)
	return n.Run()
}

func UseCustomeDNS(dns []string) {
	resolver := net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, address string) (conn net.Conn, err error) {
			for _, addr := range dns {
				if conn, err = net.Dial("udp", addr+":53"); err != nil {
					continue
				} else {
					return conn, nil
				}
			}
			return
		},
	}
	net.DefaultResolver = &resolver
}