package net

import (
	"b612.me/starlog"
	"b612.me/starnet"
	"context"
	"encoding/json"
	"fmt"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
	"net"
	"strings"
	"sync/atomic"
	"time"
)

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
}

func Traceroute(address string, dns string, maxHops int, timeout time.Duration, ipinfoAddr string) {
	ipinfo := net.ParseIP(address)
	if ipinfo == nil {
		{
			if dns != "" {
				useCustomeDNS([]string{dns})
				starlog.Infoln("使用自定义DNS服务器：", dns)
			} else {
				starlog.Infoln("使用系统默认DNS服务器")
			}
			addr, err := net.ResolveIPAddr("ip", address)
			if err != nil {
				starlog.Errorln("IP地址解析失败：", address, err)
				return
			}
			starlog.Infoln("解析IP地址：", addr.String())
			address = addr.String()
		}
	}
	traceroute(address, maxHops, timeout, ipinfoAddr)
}
func traceroute(address string, maxHops int, timeout time.Duration, ipinfoAddr string) {
	ipinfo := net.ParseIP(address)
	if ipinfo == nil {
		starlog.Errorln("IP地址解析失败：", address)
		return
	}

	var echoType icmp.Type = ipv4.ICMPTypeEcho
	var exceededType icmp.Type = ipv4.ICMPTypeTimeExceeded
	var replyType icmp.Type = ipv4.ICMPTypeEchoReply
	var proto = 1
	var network = "ip4:icmp"
	var resolveIP = "ip4"
	if ipinfo.To4() == nil {
		network = "ip6:ipv6-icmp"
		resolveIP = "ip6"
		echoType = ipv6.ICMPTypeEchoRequest
		exceededType = ipv6.ICMPTypeTimeExceeded
		replyType = ipv6.ICMPTypeEchoReply
		proto = 58
	}
	c, err := icmp.ListenPacket(network, "0.0.0.0")
	if err != nil {
		fmt.Println(err)
		return
	}
	defer c.Close()

	dst, err := net.ResolveIPAddr(resolveIP, address)
	if err != nil {
		starlog.Errorln("IP地址解析失败：", address, err)
		return
	}
	if maxHops == 0 {
		maxHops = 32
	}
	firstTargetHop := int32(maxHops + 1)
	if timeout == 0 {
		timeout = time.Second * 3
	}
exitfor:
	for i := 1; i <= maxHops; i++ {
		if atomic.LoadInt32(&firstTargetHop) <= int32(i) {
			return
		}
		m := icmp.Message{
			Type: echoType, Code: 0,
			Body: &icmp.Echo{
				ID: i, Seq: i,
				Data: []byte("B612.ME-ROUTER-TRACE"),
			},
		}

		b, err := m.Marshal(nil)
		if err != nil {
			fmt.Printf("%d\tMarshal error: %v\n", i, err)
			continue
		}

		if network == "ip4:icmp" {
			if err := c.IPv4PacketConn().SetTTL(i); err != nil {
				fmt.Printf("%d\tSetTTL error: %v\n", i, err)
				continue
			}
		} else {
			if err := c.IPv6PacketConn().SetHopLimit(i); err != nil {
				fmt.Printf("%d\tSetHopLimit error: %v\n", i, err)
				continue
			}
		}

		start := time.Now()
		n, err := c.WriteTo(b, dst)
		if err != nil {
			fmt.Printf("%d\tWriteTo error: %v\n", i, err)
			continue
		} else if n != len(b) {
			fmt.Printf("%d\tWrite Short: %v Expected: %v\n", i, n, len(b))
			continue
		}

		reply := make([]byte, 1500)
		err = c.SetReadDeadline(time.Now().Add(timeout))
		if err != nil {
			fmt.Printf("%d\tSetReadDeadline error: %v\n", i, err)
			continue
		}
		n, peer, err := c.ReadFrom(reply)
		if err != nil {
			fmt.Printf("%d\tReadFrom error: %v\n", i, err)
			continue
		}
		duration := time.Since(start)

		rm, err := icmp.ParseMessage(proto, reply[:n])
		if err != nil {
			fmt.Printf("%d\tParseMessage error: %v\n", i, err)
			return
		}

		switch rm.Type {
		case exceededType:
			fmt.Printf("%d\thops away:\t%s\t(%s) %s\n", i, peer, duration, GetIPInfo(peer.String(), ipinfoAddr))
		case replyType:
			fmt.Printf("%d\thops away:\t%s\t(%s) %s\n", i, peer, duration, GetIPInfo(peer.String(), ipinfoAddr))
			break exitfor
		default:
			fmt.Printf("%d\tgot %+v from %v; want echo reply;%s\n", i, rm, peer, GetIPInfo(peer.String(), ipinfoAddr))
		}
	}
}

func GetIPInfo(ip string, addr string) string {
	if addr == "" {
		return ""
	}
	uri := strings.ReplaceAll(addr, "{ip}", ip)
	res, err := starnet.Curl(starnet.NewRequests(uri, nil, "GET", starnet.WithTimeout(time.Second*2), starnet.WithDialTimeout(time.Second*3)))
	if err != nil {
		return "获取IP信息失败：" + err.Error()
	}
	var ipinfo IPInfo
	err = json.Unmarshal(res.RecvData, &ipinfo)
	if err != nil {
		return "解析IP信息失败：" + err.Error()
	}
	return fmt.Sprintf("%s %s %s %s %s", ipinfo.CountryName, ipinfo.RegionName, ipinfo.CityName, ipinfo.OwnerDomain, ipinfo.ISP)
}

type IPInfo struct {
	CountryName string `json:"country_name"`
	RegionName  string `json:"region_name"`
	CityName    string `json:"city_name"`
	OwnerDomain string `json:"owner_domain"`
	Ip          string `json:"ip"`
	ISP         string `json:"isp_domain"`
	Err         string `json:"err"`
}