package net

import (
	"b612.me/starlog"
	"b612.me/starnet"
	"context"
	"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, bindaddr string, dns string, maxHops int, timeout time.Duration, ipinfoAddr string, hideIncorrect bool) {
	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, bindaddr, maxHops, timeout, ipinfoAddr, hideIncorrect)
}

func traceroute(address string, bindaddr string, maxHops int, timeout time.Duration, ipinfoAddr string, hideIncorrect bool) {
	ipinfo := net.ParseIP(address)
	if ipinfo == nil {
		starlog.Errorln("IP地址解析失败：", address)
		return
	}

	var (
		echoType     icmp.Type
		exceededType icmp.Type
		replyType    icmp.Type
		unreachType  icmp.Type
		proto        int
		network      string
		resolveIP    string
		isIPv4       bool
	)

	if ipinfo.To4() != nil {
		echoType = ipv4.ICMPTypeEcho
		exceededType = ipv4.ICMPTypeTimeExceeded
		replyType = ipv4.ICMPTypeEchoReply
		unreachType = ipv4.ICMPTypeDestinationUnreachable
		proto = 1
		network = "ip4:icmp"
		resolveIP = "ip4"
		isIPv4 = true
	} else {
		echoType = ipv6.ICMPTypeEchoRequest
		exceededType = ipv6.ICMPTypeTimeExceeded
		replyType = ipv6.ICMPTypeEchoReply
		unreachType = ipv6.ICMPTypeDestinationUnreachable
		proto = 58
		network = "ip6:ipv6-icmp"
		resolveIP = "ip6"
		isIPv4 = false
	}
	if bindaddr == "" {
		bindaddr = "0.0.0.0"
		if !isIPv4 {
			bindaddr = "::"
		}
	}

	c, err := icmp.ListenPacket(network, bindaddr)
	if err != nil {
		starlog.Errorln("监听失败:", err)
		return
	}
	defer c.Close()

	if maxHops == 0 {
		maxHops = 32
	}
	firstTargetHop := int32(maxHops + 1)
	if timeout == 0 {
		timeout = time.Second * 3
	}

exitfor:
	for ttl := 1; ttl <= maxHops; ttl++ {
		if atomic.LoadInt32(&firstTargetHop) <= int32(ttl) {
			return
		}

		dst, err := net.ResolveIPAddr(resolveIP, address)
		if err != nil {
			starlog.Errorln("解析失败:", address, err)
			return
		}

		// 构造ICMP报文
		msg := icmp.Message{
			Type: echoType, Code: 0,
			Body: &icmp.Echo{
				ID:   ttl, // 使用TTL作为ID
				Seq:  ttl, // 使用TTL作为序列号
				Data: []byte("B612.ME-ROUTER-TRACE"),
			},
		}

		msgBytes, err := msg.Marshal(nil)
		if err != nil {
			starlog.Warningf("%d\t封包失败: %v\n", ttl, err)
			continue
		}

		// 设置TTL/HopLimit
		if network == "ip4:icmp" {
			if err := c.IPv4PacketConn().SetTTL(ttl); err != nil {
				starlog.Warningf("%d\t设置TTL失败: %v\n", ttl, err)
				continue
			}
		} else {
			if err := c.IPv6PacketConn().SetHopLimit(ttl); err != nil {
				starlog.Warningf("%d\t设置HopLimit失败: %v\n", ttl, err)
				continue
			}
		}

		startTime := time.Now()
		if _, err := c.WriteTo(msgBytes, dst); err != nil {
			starlog.Warningf("%d\t发送失败: %v\n", ttl, err)
			continue
		}

		// 接收响应处理
		timeoutCh := time.After(timeout)
		responsesReceived := 0

	recvLoop:
		for responsesReceived < 3 {
			select {
			case <-timeoutCh:
				if responsesReceived == 0 {
					fmt.Printf("%d\t*\n", ttl)
				}
				break recvLoop
			default:
				reply := make([]byte, 1500)
				if err := c.SetReadDeadline(time.Now().Add(50 * time.Millisecond)); err != nil {
					break recvLoop
				}

				n, peer, err := c.ReadFrom(reply)
				if err != nil {
					if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
						continue
					}
					starlog.Debugf("%d\t接收错误: %v", ttl, err)
					continue
				}

				// 解析响应
				rm, err := icmp.ParseMessage(proto, reply[:n])
				if err != nil {
					starlog.Debugf("%d\t解析错误: %v", ttl, err)
					continue
				}

				// 验证响应匹配
				if match := checkResponseMatch(rm, ttl, peer.String() == dst.String(), isIPv4, exceededType, replyType, unreachType); match {
					duration := time.Since(startTime)
					fmt.Printf("%d\t%s\t%s\t%s\n",
						ttl,
						peer,
						duration.Round(time.Millisecond),
						GetIPInfo(peer.String(), ipinfoAddr),
					)
					responsesReceived++

					if peer.String() == dst.String() {
						atomic.StoreInt32(&firstTargetHop, int32(ttl))
						break exitfor
					}
				}
			}
		}
	}
}

func checkResponseMatch(rm *icmp.Message, ttl int, isFinal bool, isIPv4 bool,
	exceededType, replyType, unreachType icmp.Type) bool {

	switch {
	case rm.Type == exceededType:
		if body, ok := rm.Body.(*icmp.TimeExceeded); ok {
			return validateOriginalPacket(body.Data, ttl, isIPv4)
		}

	case rm.Type == replyType:
		if isFinal {
			if body, ok := rm.Body.(*icmp.Echo); ok {
				return body.ID == ttl && body.Seq == ttl
			}
		}
		return false

	case rm.Type == unreachType:
		if body, ok := rm.Body.(*icmp.DstUnreach); ok {
			return validateOriginalPacket(body.Data, ttl, isIPv4)
		}
	}
	return false
}

func validateOriginalPacket(data []byte, ttl int, isIPv4 bool) bool {
	var (
		proto  byte
		header []byte
	)

	if isIPv4 {
		if len(data) < 20+8 {
			return false
		}
		ihl := data[0] & 0x0F
		if ihl < 5 {
			return false
		}
		proto = data[9]
		header = data[:ihl*4]
	} else {
		if len(data) < 40+8 {
			return false
		}
		proto = data[6]
		header = data[:40]
	}

	if proto != 1 && proto != 58 {
		return false
	}

	payload := data[len(header):]
	if len(payload) < 8 {
		return false
	}

	if isIPv4 {
		return payload[0] == 8 && // ICMP Echo Request
			payload[4] == byte(ttl>>8) &&
			payload[5] == byte(ttl) &&
			payload[6] == byte(ttl>>8) &&
			payload[7] == byte(ttl)
	} else {
		return payload[0] == 128 && // ICMPv6 Echo Request
			payload[4] == byte(ttl>>8) &&
			payload[5] == byte(ttl) &&
			payload[6] == byte(ttl>>8) &&
			payload[7] == byte(ttl)
	}
}

func GetIPInfo(ip string, addr string) string {
	if addr == "" {
		return ""
	}
	uri := strings.ReplaceAll(addr, "{ip}", ip)
	res, err := starnet.Curl(starnet.NewSimpleRequest(uri, "GET",
		starnet.WithTimeout(time.Second*2),
		starnet.WithDialTimeout(time.Second*3)))

	if err != nil {
		return "IP信息获取失败"
	}

	var ipinfo IPInfo
	if err := res.Body().Unmarshal(&ipinfo); err != nil {
		return "IP信息解析失败"
	}

	return fmt.Sprintf("%s %s %s",
		ipinfo.CountryName,
		ipinfo.RegionName,
		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"`
}