package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

type BroadcastMessage struct {
	Type string `json:"type"`
	IP   string `json:"ip"`
	Port int    `json:"port"`
	Name string `json:"name"`
}

func getGatewayIP() (string, error) {
	var cmd *exec.Cmd
	if runtime.GOOS == "linux" {
		cmd = exec.Command("sh", "-c", "ip route | grep default | awk '{print $3}'")
	} else if runtime.GOOS == "darwin" {
		cmd = exec.Command("sh", "-c", "netstat -rn | grep default | awk '{print $2}'")
	} else {
		return "", fmt.Errorf("unsupported OS")
	}

	out, err := cmd.Output()
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

func ipInSameSubnet(ip1 net.IP, ip2 net.IP, mask net.IPMask) bool {
	if mask == nil {
		return false
	}
	net1 := ip1.Mask(mask)
	net2 := ip2.Mask(mask)
	return net1.Equal(net2)
}

func getLocalIP() string {
	gatewayStr, err := getGatewayIP()
	if err != nil {
		fmt.Println("Failed to get gateway IP:", err)
		return "127.0.0.1"
	}
	gateway := net.ParseIP(gatewayStr)
	if gateway == nil {
		fmt.Println("Invalid gateway IP")
		return "127.0.0.1"
	}

	interfaces, err := net.Interfaces()
	if err != nil {
		fmt.Println("Failed to get network interfaces:", err)
		return "127.0.0.1"
	}

	for _, iface := range interfaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface down
		}
		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}
		for _, addr := range addrs {
			var ip net.IP
			var mask net.IPMask

			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
				mask = v.Mask
			case *net.IPAddr:
				ip = v.IP
			}

			if ip == nil || ip.IsLoopback() || ip.To4() == nil {
				continue
			}

			if ipInSameSubnet(ip, gateway, mask) {
				return ip.String()
			}
		}
	}
	return "127.0.0.1"
}

func main() {
	// 参数
	var (
		port     = flag.Int("port", 8000, "Device service port")
		interval = flag.Int("interval", 2, "Broadcast interval in seconds")
		name     = flag.String("name", "My-AI-Server", "Device name")
	)
	flag.Parse()

	ip := getLocalIP()
	if ip == "127.0.0.1" {
		fmt.Println("Warning: failed to find suitable IP address, using loopback")
	}

	addr := net.UDPAddr{
		IP:   net.IPv4bcast,
		Port: 9876,
	}
	conn, err := net.DialUDP("udp4", nil, &addr)
	if err != nil {
		fmt.Println("Failed to dial UDP broadcast:", err)
		os.Exit(1)
	}
	defer conn.Close()

	for {
		msg := BroadcastMessage{
			Type: "ai_server_announce",
			IP:   ip,
			Port: *port,
			Name: *name,
		}
		data, _ := json.Marshal(msg)
		_, err := conn.Write(data)
		if err != nil {
			fmt.Println("Broadcast failed:", err)
		} else {
			fmt.Printf("Broadcasted: %s:%d (%s)\n", ip, *port, *name)
		}
		time.Sleep(time.Duration(*interval) * time.Second)
	}
}