package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
)

// ============ 通用数据结构 ============
type TestPacket struct {
	SeqNum    uint64
	Timestamp int64
	Type      string // "tcp", "udp", "ping"
	Data      []byte
}

type Stats struct {
	PacketsReceived uint64
	PacketsLost     uint64
	LastSeqNum      uint64
	RTTSamples      []time.Duration
}

type Metrics struct {
	PacketsSent     uint64
	PacketsReceived uint64
	PacketsLost     uint64
	RTTSamples      []time.Duration
	MinRTT          time.Duration
	MaxRTT          time.Duration
	AvgRTT          time.Duration
	Jitter          time.Duration
}

type TestReport struct {
	Timestamp      time.Time
	TestDuration   time.Duration
	TargetHost     string
	TCPMetrics     *Metrics
	UDPMetrics     *Metrics
	TracerouteHops []HopInfo
}

type HopInfo struct {
	TTL     int
	Address string
	RTT     time.Duration
}

// ============ 服务端实现 ============
type NetworkServer struct {
	tcpAddr   string
	udpAddr   string
	tcpStats  *Stats
	udpStats  *Stats
	statsLock sync.RWMutex
}

func NewNetworkServer(tcpPort, udpPort int) *NetworkServer {
	return &NetworkServer{
		tcpAddr:  fmt.Sprintf(":%d", tcpPort),
		udpAddr:  fmt.Sprintf(":%d", udpPort),
		tcpStats: &Stats{},
		udpStats: &Stats{},
	}
}

func (ns *NetworkServer) Start() {
	log.Printf("========== 网络质量检测服务端 ==========")
	log.Printf("TCP监听端口: %s", ns.tcpAddr)
	log.Printf("UDP监听端口: %s", ns.udpAddr)
	log.Printf("服务器已启动，等待客户端连接...")
	log.Printf("========================================\n")

	// 启动TCP服务器
	go ns.serveTCP()
	// 启动UDP服务器
	go ns.serveUDP()

	select {}
}

func (ns *NetworkServer) serveTCP() {
	listener, err := net.Listen("tcp", ns.tcpAddr)
	if err != nil {
		log.Fatalf("TCP监听失败: %v", err)
	}
	defer listener.Close()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("TCP连接接受错误: %v", err)
			continue
		}
		go ns.handleTCPConnection(conn)
	}
}

func (ns *NetworkServer) handleTCPConnection(conn net.Conn) {
	defer conn.Close()
	log.Printf("[TCP] 新连接来自 %s", conn.RemoteAddr())

	buf := make([]byte, 8192)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			log.Printf("[TCP] 连接 %s 断开", conn.RemoteAddr())
			return
		}

		var packet TestPacket
		if err := json.Unmarshal(buf[:n], &packet); err != nil {
			continue
		}

		receiveTime := time.Now().UnixNano()
		ns.updateStats(ns.tcpStats, packet.SeqNum)

		// 立即回显数据包
		response := TestPacket{
			SeqNum:    packet.SeqNum,
			Timestamp: receiveTime,
			Type:      "tcp_response",
			Data:      packet.Data,
		}

		data, _ := json.Marshal(response)
		conn.Write(data)
	}
}

func (ns *NetworkServer) serveUDP() {
	addr, err := net.ResolveUDPAddr("udp", ns.udpAddr)
	if err != nil {
		log.Fatalf("UDP地址解析失败: %v", err)
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Fatalf("UDP监听失败: %v", err)
	}
	defer conn.Close()

	log.Printf("[UDP] 监听在 %s", ns.udpAddr)

	buf := make([]byte, 8192)
	for {
		n, remoteAddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			log.Printf("UDP读取错误: %v", err)
			continue
		}

		var packet TestPacket
		if err := json.Unmarshal(buf[:n], &packet); err != nil {
			continue
		}

		receiveTime := time.Now().UnixNano()
		ns.updateStats(ns.udpStats, packet.SeqNum)

		// 回显UDP数据包
		response := TestPacket{
			SeqNum:    packet.SeqNum,
			Timestamp: receiveTime,
			Type:      "udp_response",
			Data:      packet.Data,
		}

		data, _ := json.Marshal(response)
		conn.WriteToUDP(data, remoteAddr)
	}
}

func (ns *NetworkServer) updateStats(stats *Stats, seqNum uint64) {
	ns.statsLock.Lock()
	defer ns.statsLock.Unlock()

	stats.PacketsReceived++

	if seqNum > stats.LastSeqNum+1 {
		stats.PacketsLost += seqNum - stats.LastSeqNum - 1
	}
	stats.LastSeqNum = seqNum
}

// ============ 客户端实现 ============
type NetworkClient struct {
	targetHost   string
	tcpPort      int
	udpPort      int
	testDuration time.Duration
	packetSize   int
	reportFile   string

	tcpMetrics *Metrics
	udpMetrics *Metrics
	mu         sync.Mutex
}

func NewNetworkClient(host string, tcpPort, udpPort int, duration time.Duration) *NetworkClient {
	return &NetworkClient{
		targetHost:   host,
		tcpPort:      tcpPort,
		udpPort:      udpPort,
		testDuration: duration,
		packetSize:   1024,
		reportFile:   "network_quality_report.json",
		tcpMetrics:   &Metrics{MinRTT: time.Hour},
		udpMetrics:   &Metrics{MinRTT: time.Hour},
	}
}

func (nc *NetworkClient) testTCPLatency() error {
	addr := fmt.Sprintf("%s:%d", nc.targetHost, nc.tcpPort)
	conn, err := net.DialTimeout("tcp", addr, 5*time.Second)
	if err != nil {
		return fmt.Errorf("TCP连接失败: %v", err)
	}
	defer conn.Close()

	log.Printf("[TCP] 开始延迟测试 -> %s", addr)

	var seqNum uint64
	deadline := time.Now().Add(nc.testDuration)

	for time.Now().Before(deadline) {
		seqNum++

		packet := TestPacket{
			SeqNum:    seqNum,
			Timestamp: time.Now().UnixNano(),
			Type:      "tcp_probe",
			Data:      make([]byte, nc.packetSize),
		}

		sendTime := time.Now()

		data, _ := json.Marshal(packet)
		if _, err := conn.Write(data); err != nil {
			nc.mu.Lock()
			nc.tcpMetrics.PacketsLost++
			nc.mu.Unlock()
			continue
		}
		nc.mu.Lock()
		nc.tcpMetrics.PacketsSent++
		nc.mu.Unlock()

		buf := make([]byte, 8192)
		conn.SetReadDeadline(time.Now().Add(2 * time.Second))
		_, err := conn.Read(buf)
		if err != nil {
			nc.mu.Lock()
			nc.tcpMetrics.PacketsLost++
			nc.mu.Unlock()
			continue
		}

		rtt := time.Since(sendTime)
		nc.updateTCPMetrics(rtt)

		time.Sleep(100 * time.Millisecond)
	}

	log.Printf("[TCP] 测试完成")
	return nil
}

func (nc *NetworkClient) testUDPLatency() error {
	addr := fmt.Sprintf("%s:%d", nc.targetHost, nc.udpPort)
	raddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return fmt.Errorf("UDP地址解析失败: %v", err)
	}

	conn, err := net.DialUDP("udp", nil, raddr)
	if err != nil {
		return fmt.Errorf("UDP连接失败: %v", err)
	}
	defer conn.Close()

	log.Printf("[UDP] 开始延迟测试 -> %s", addr)

	var seqNum uint64
	deadline := time.Now().Add(nc.testDuration)
	sentPackets := make(map[uint64]time.Time)
	var wg sync.WaitGroup

	// 发送协程
	wg.Add(1)
	go func() {
		defer wg.Done()
		for time.Now().Before(deadline) {
			seqNum++

			packet := TestPacket{
				SeqNum:    seqNum,
				Timestamp: time.Now().UnixNano(),
				Type:      "udp_probe",
				Data:      make([]byte, nc.packetSize),
			}

			sendTime := time.Now()
			nc.mu.Lock()
			sentPackets[seqNum] = sendTime
			nc.mu.Unlock()

			data, _ := json.Marshal(packet)
			conn.Write(data)

			nc.mu.Lock()
			nc.udpMetrics.PacketsSent++
			nc.mu.Unlock()

			time.Sleep(100 * time.Millisecond)
		}
	}()

	// 接收协程
	buf := make([]byte, 8192)
	for time.Now().Before(deadline.Add(3 * time.Second)) {
		conn.SetReadDeadline(time.Now().Add(1 * time.Second))
		n, err := conn.Read(buf)
		if err != nil {
			continue
		}

		var response TestPacket
		if err := json.Unmarshal(buf[:n], &response); err != nil {
			continue
		}

		nc.mu.Lock()
		if sendTime, ok := sentPackets[response.SeqNum]; ok {
			rtt := time.Since(sendTime)
			nc.updateUDPMetrics(rtt)
			delete(sentPackets, response.SeqNum)
		}
		nc.mu.Unlock()
	}

	wg.Wait()

	nc.mu.Lock()
	nc.udpMetrics.PacketsLost = uint64(len(sentPackets))
	nc.mu.Unlock()

	log.Printf("[UDP] 测试完成")
	return nil
}

func (nc *NetworkClient) performTraceroute() ([]HopInfo, error) {
	log.Printf("[Traceroute] 路由追踪到 %s", nc.targetHost)

	hops := make([]HopInfo, 0, 30)
	maxTTL := 30
	timeout := 2 * time.Second

	for ttl := 1; ttl <= maxTTL; ttl++ {
		hopInfo, reached, err := nc.probeHop(ttl, timeout)
		if err != nil {
			continue
		}

		hops = append(hops, hopInfo)
		if hopInfo.Address != "*" {
			log.Printf("  %2d  %-15s  %.2fms", ttl, hopInfo.Address,
				float64(hopInfo.RTT.Microseconds())/1000.0)
		} else {
			log.Printf("  %2d  *", ttl)
		}

		if reached {
			break
		}
	}

	return hops, nil
}

func (nc *NetworkClient) probeHop(ttl int, timeout time.Duration) (HopInfo, bool, error) {
	conn, err := icmp.ListenPacket("ip4:icmp", "0.0.0.0")
	if err != nil {
		return HopInfo{}, false, err
	}
	defer conn.Close()

	if err := conn.IPv4PacketConn().SetTTL(ttl); err != nil {
		return HopInfo{}, false, err
	}

	msg := icmp.Message{
		Type: ipv4.ICMPTypeEcho,
		Code: 0,
		Body: &icmp.Echo{
			ID:   os.Getpid() & 0xffff,
			Seq:  ttl,
			Data: []byte("TRACEROUTE"),
		},
	}

	msgBytes, err := msg.Marshal(nil)
	if err != nil {
		return HopInfo{}, false, err
	}

	dst, err := net.ResolveIPAddr("ip4", nc.targetHost)
	if err != nil {
		return HopInfo{}, false, err
	}

	start := time.Now()
	if _, err := conn.WriteTo(msgBytes, dst); err != nil {
		return HopInfo{}, false, err
	}

	reply := make([]byte, 1500)
	conn.SetReadDeadline(time.Now().Add(timeout))
	_, peer, err := conn.ReadFrom(reply)
	rtt := time.Since(start)

	if err != nil {
		return HopInfo{TTL: ttl, Address: "*", RTT: 0}, false, nil
	}

	hopAddr := peer.String()
	reachedTarget := (hopAddr == dst.String())

	return HopInfo{
		TTL:     ttl,
		Address: hopAddr,
		RTT:     rtt,
	}, reachedTarget, nil
}

func (nc *NetworkClient) updateTCPMetrics(rtt time.Duration) {
	nc.mu.Lock()
	defer nc.mu.Unlock()

	nc.tcpMetrics.PacketsReceived++
	nc.tcpMetrics.RTTSamples = append(nc.tcpMetrics.RTTSamples, rtt)

	if rtt < nc.tcpMetrics.MinRTT {
		nc.tcpMetrics.MinRTT = rtt
	}
	if rtt > nc.tcpMetrics.MaxRTT {
		nc.tcpMetrics.MaxRTT = rtt
	}
}

func (nc *NetworkClient) updateUDPMetrics(rtt time.Duration) {
	nc.udpMetrics.PacketsReceived++
	nc.udpMetrics.RTTSamples = append(nc.udpMetrics.RTTSamples, rtt)

	if rtt < nc.udpMetrics.MinRTT {
		nc.udpMetrics.MinRTT = rtt
	}
	if rtt > nc.udpMetrics.MaxRTT {
		nc.udpMetrics.MaxRTT = rtt
	}
}

func (nc *NetworkClient) calculateMetrics() {
	// 计算TCP平均RTT和抖动
	if len(nc.tcpMetrics.RTTSamples) > 0 {
		var sum time.Duration
		for _, rtt := range nc.tcpMetrics.RTTSamples {
			sum += rtt
		}
		nc.tcpMetrics.AvgRTT = sum / time.Duration(len(nc.tcpMetrics.RTTSamples))

		var jitterSum time.Duration
		for i := 1; i < len(nc.tcpMetrics.RTTSamples); i++ {
			diff := nc.tcpMetrics.RTTSamples[i] - nc.tcpMetrics.RTTSamples[i-1]
			if diff < 0 {
				diff = -diff
			}
			jitterSum += diff
		}
		if len(nc.tcpMetrics.RTTSamples) > 1 {
			nc.tcpMetrics.Jitter = jitterSum / time.Duration(len(nc.tcpMetrics.RTTSamples)-1)
		}
	}

	// 计算UDP平均RTT和抖动
	if len(nc.udpMetrics.RTTSamples) > 0 {
		var sum time.Duration
		for _, rtt := range nc.udpMetrics.RTTSamples {
			sum += rtt
		}
		nc.udpMetrics.AvgRTT = sum / time.Duration(len(nc.udpMetrics.RTTSamples))

		var jitterSum time.Duration
		for i := 1; i < len(nc.udpMetrics.RTTSamples); i++ {
			diff := nc.udpMetrics.RTTSamples[i] - nc.udpMetrics.RTTSamples[i-1]
			if diff < 0 {
				diff = -diff
			}
			jitterSum += diff
		}
		if len(nc.udpMetrics.RTTSamples) > 1 {
			nc.udpMetrics.Jitter = jitterSum / time.Duration(len(nc.udpMetrics.RTTSamples)-1)
		}
	}
}

func (nc *NetworkClient) generateReport(hops []HopInfo) error {
	nc.calculateMetrics()

	report := TestReport{
		Timestamp:      time.Now(),
		TestDuration:   nc.testDuration,
		TargetHost:     nc.targetHost,
		TCPMetrics:     nc.tcpMetrics,
		UDPMetrics:     nc.udpMetrics,
		TracerouteHops: hops,
	}

	// 将报告序列化为JSON（单行，不格式化）
	data, err := json.Marshal(report)
	if err != nil {
		return err
	}

	// 以追加模式打开文件，如果不存在则创建
	file, err := os.OpenFile(nc.reportFile, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		return fmt.Errorf("打开报告文件失败: %v", err)
	}
	defer file.Close()

	// 写入JSON数据并添加换行符（JSON Lines格式）
	if _, err := file.Write(data); err != nil {
		return fmt.Errorf("写入报告数据失败: %v", err)
	}
	if _, err := file.WriteString("\n"); err != nil {
		return fmt.Errorf("写入换行符失败: %v", err)
	}

	log.Printf("测试报告已追加至: %s", nc.reportFile)

	nc.printReport(&report)

	return nil
}

func (nc *NetworkClient) printReport(report *TestReport) {
	fmt.Println("\n" + strings.Repeat("=", 50))
	fmt.Println("           网络质量检测报告")
	fmt.Println(strings.Repeat("=", 50))
	fmt.Printf("测试时间: %s\n", report.Timestamp.Format("2006-01-02 15:04:05"))
	fmt.Printf("目标主机: %s\n", report.TargetHost)
	fmt.Printf("测试时长: %v\n", report.TestDuration)
	fmt.Println(strings.Repeat("-", 50))

	fmt.Println("\n【TCP 测试结果】")
	fmt.Printf("  发送包数: %d\n", report.TCPMetrics.PacketsSent)
	fmt.Printf("  接收包数: %d\n", report.TCPMetrics.PacketsReceived)
	if report.TCPMetrics.PacketsSent > 0 {
		lossRate := float64(report.TCPMetrics.PacketsLost) / float64(report.TCPMetrics.PacketsSent) * 100
		fmt.Printf("  丢包数量: %d (丢包率: %.2f%%)\n", report.TCPMetrics.PacketsLost, lossRate)
	}
	if report.TCPMetrics.MinRTT < time.Hour {
		fmt.Printf("  最小RTT:  %v\n", report.TCPMetrics.MinRTT)
		fmt.Printf("  平均RTT:  %v\n", report.TCPMetrics.AvgRTT)
		fmt.Printf("  最大RTT:  %v\n", report.TCPMetrics.MaxRTT)
		fmt.Printf("  抖动:     %v\n", report.TCPMetrics.Jitter)
	}

	fmt.Println("\n【UDP 测试结果】")
	fmt.Printf("  发送包数: %d\n", report.UDPMetrics.PacketsSent)
	fmt.Printf("  接收包数: %d\n", report.UDPMetrics.PacketsReceived)
	if report.UDPMetrics.PacketsSent > 0 {
		lossRate := float64(report.UDPMetrics.PacketsLost) / float64(report.UDPMetrics.PacketsSent) * 100
		fmt.Printf("  丢包数量: %d (丢包率: %.2f%%)\n", report.UDPMetrics.PacketsLost, lossRate)
	}
	if report.UDPMetrics.MinRTT < time.Hour {
		fmt.Printf("  最小RTT:  %v\n", report.UDPMetrics.MinRTT)
		fmt.Printf("  平均RTT:  %v\n", report.UDPMetrics.AvgRTT)
		fmt.Printf("  最大RTT:  %v\n", report.UDPMetrics.MaxRTT)
		fmt.Printf("  抖动:     %v\n", report.UDPMetrics.Jitter)
	}

	fmt.Println(strings.Repeat("=", 50))
	fmt.Printf("报告已保存至: %s\n\n", nc.reportFile)
}

func (nc *NetworkClient) RunScheduledTests(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		log.Printf("\n========== 开始定时测试 [%s] ==========",
			time.Now().Format("2006-01-02 15:04:05"))

		// 执行TCP测试
		if err := nc.testTCPLatency(); err != nil {
			log.Printf("TCP测试错误: %v", err)
		}

		// 重置UDP指标
		nc.mu.Lock()
		nc.udpMetrics = &Metrics{MinRTT: time.Hour}
		nc.mu.Unlock()

		// 执行UDP测试
		if err := nc.testUDPLatency(); err != nil {
			log.Printf("UDP测试错误: %v", err)
		}

		// 执行Traceroute
		hops, err := nc.performTraceroute()
		if err != nil {
			log.Printf("Traceroute错误: %v", err)
		}

		// 生成报告
		if err := nc.generateReport(hops); err != nil {
			log.Printf("报告生成错误: %v", err)
		}

		// 重置TCP指标准备下一轮
		nc.mu.Lock()
		nc.tcpMetrics = &Metrics{MinRTT: time.Hour}
		nc.mu.Unlock()

		log.Printf("========== 测试完成，等待下一轮 ==========\n")

		<-ticker.C
	}
}

// ============ 主程序 ============
func main() {
	// 定义命令行参数
	var (
		mode         = flag.String("mode", "", "运行模式: server(服务端) 或 client(客户端)")
		tcpPort      = flag.Int("tcp", 9001, "TCP测试端口")
		udpPort      = flag.Int("udp", 9002, "UDP测试端口")
		targetHost   = flag.String("target", "", "目标主机地址(客户端模式必需)")
		testDuration = flag.Int("duration", 60, "单次测试时长(秒)")
		interval     = flag.Int("interval", 3600, "定时测试间隔(秒), 0表示只执行一次")
	)

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "\n网络质量检测工具\n\n")
		fmt.Fprintf(os.Stderr, "用法:\n")
		fmt.Fprintf(os.Stderr, "  服务端模式: %s -mode server [-tcp 端口] [-udp 端口]\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  客户端模式: %s -mode client -target 目标IP [-tcp 端口] [-udp 端口] [-duration 秒] [-interval 秒]\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "参数说明:\n")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n示例:\n")
		fmt.Fprintf(os.Stderr, "  # 启动服务端，使用默认端口\n")
		fmt.Fprintf(os.Stderr, "  %s -mode server\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  # 启动服务端，自定义端口\n")
		fmt.Fprintf(os.Stderr, "  %s -mode server -tcp 8001 -udp 8002\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  # 启动客户端，单次测试60秒\n")
		fmt.Fprintf(os.Stderr, "  %s -mode client -target 192.168.1.100 -duration 60 -interval 0\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  # 启动客户端，每小时测试一次\n")
		fmt.Fprintf(os.Stderr, "  %s -mode client -target 192.168.1.100 -interval 3600\n\n", os.Args[0])
	}

	flag.Parse()

	// 验证参数
	if *mode == "" {
		fmt.Fprintf(os.Stderr, "错误: 必须指定运行模式 -mode server 或 -mode client\n\n")
		flag.Usage()
		os.Exit(1)
	}

	switch *mode {
	case "server":
		// 服务端模式
		server := NewNetworkServer(*tcpPort, *udpPort)
		server.Start()

	case "client":
		// 客户端模式
		if *targetHost == "" {
			fmt.Fprintf(os.Stderr, "错误: 客户端模式必须指定 -target 参数\n\n")
			flag.Usage()
			os.Exit(1)
		}

		client := NewNetworkClient(*targetHost, *tcpPort, *udpPort, time.Duration(*testDuration)*time.Second)

		if *interval == 0 {
			// 单次执行
			log.Printf("开始单次网络质量测试...")

			if err := client.testTCPLatency(); err != nil {
				log.Printf("TCP测试错误: %v", err)
			}

			if err := client.testUDPLatency(); err != nil {
				log.Printf("UDP测试错误: %v", err)
			}

			hops, err := client.performTraceroute()
			if err != nil {
				log.Printf("Traceroute错误: %v", err)
			}

			if err := client.generateReport(hops); err != nil {
				log.Printf("报告生成错误: %v", err)
			}
		} else {
			// 定时执行
			log.Printf("开始定时网络质量监控，间隔: %d秒", *interval)
			client.RunScheduledTests(time.Duration(*interval) * time.Second)
		}

	default:
		fmt.Fprintf(os.Stderr, "错误: 无效的运行模式 '%s'，必须是 'server' 或 'client'\n\n", *mode)
		flag.Usage()
		os.Exit(1)
	}
}