package utils

import (
    "fmt"
    "log"
    "net"
    "os"
    "os/exec"
    "runtime"
    "sort"
    "strings"
    "sync"
    "time"
	"testing"
)

// DeviceInfo 存储设备信息
type DeviceInfo struct {
    IP      string
    MAC     string
    Hostname string
    Vendor  string
}

// MAC厂商前缀映射（部分常见厂商）
var macVendors = map[string]string{
    "00:0C:29": "VMware",
    "00:50:56": "VMware",
    "00:1A:2B": "Cisco",
    "00:1C:42": "Apple",
    "00:1D:4F": "Apple",
    "00:23:DF": "Apple",
    "00:25:BC": "Apple",
    "08:00:27": "VirtualBox",
    "00:1B:21": "Intel",
    "00:21:5A": "Intel",
    "00:22:FA": "Intel",
    "B8:27:EB": "Raspberry Pi",
    "DC:A6:32": "Raspberry Pi",
    "28:16:AD": "TP-Link",
    "14:CC:20": "TP-Link",
    "B0:95:8E": "TP-Link",
    "C8:3A:35": "Tenda",
    "80:89:17": "Tenda",
    "A0:F3:C1": "TP-Link",
    "D4:61:DA": "TP-Link",
    "C0:56:27": "TP-Link",
    "F8:1A:67": "TP-Link",
    "F8:1A:2B": "TP-Link",
    "E4:8B:7F": "HUAWEI",
    "00:9A:CD": "HUAWEI",
    "34:29:8F": "HUAWEI",
    "80:E6:50": "HUAWEI",
    "B4:0B:44": "Xiaomi",
    "E4:46:DA": "Xiaomi",
    "0C:1D:AF": "Xiaomi",
    "D8:63:75": "Xiaomi",
    "D0:C7:C0": "Xiaomi",
    "FC:64:BA": "Samsung",
    "5C:49:79": "Samsung",
    "D0:33:11": "Samsung",
    "38:48:4C": "Samsung",
    "88:36:6C": "Samsung",
    "00:1E:65": "Samsung",
    "C0:BD:D1": "D-Link",
    "1C:AF:05": "D-Link",
    "BC:F6:85": "D-Link",
}

// getLocalIP 获取本地IP地址
func getLocalIP() (net.IP, *net.IPNet) {
    conn, err := net.Dial("udp", "8.8.8.8:80")
    if err != nil {
        log.Fatal(err)
    }
    defer conn.Close()
    
    localAddr := conn.LocalAddr().(*net.UDPAddr)
    
    interfaces, err := net.Interfaces()
    if err != nil {
        log.Fatal(err)
    }
    
    for _, iface := range interfaces {
        addrs, err := iface.Addrs()
        if err != nil {
            continue
        }
        
        for _, addr := range addrs {
            switch v := addr.(type) {
            case *net.IPNet:
                if v.IP.To4() != nil && v.IP.Equal(localAddr.IP) {
                    return v.IP, v
                }
            }
        }
    }
    
    return nil, nil
}

// getIPRange 获取IP范围
func getIPRange(ipNet *net.IPNet) ([]string, error) {
    var ips []string
    
    // 获取网络地址和掩码
    mask := ipNet.Mask
    network := ipNet.IP.To4()
    
    if network == nil {
        return nil, fmt.Errorf("IPv6 not supported")
    }
    
    // 计算网络大小
    ones, bits := mask.Size()
    totalIPs := 1 << (bits - ones)
    
    // 排除网络地址和广播地址
    start := 1
    end := totalIPs - 2
    
    if totalIPs <= 2 { // /31 或 /32 网络
        start = 0
        end = totalIPs - 1
    }
    
    // 生成IP列表
    for i := start; i <= end; i++ {
        ip := make(net.IP, len(network))
        copy(ip, network)
        
        // 计算当前IP
        for j := len(ip) - 1; j >= 0; j-- {
            ip[j] += byte(i >> uint((len(ip)-1-j)*8))
        }
        
        ips = append(ips, ip.String())
    }
    
    return ips, nil
}

// pingIP 使用ping检测主机是否在线
func pingIP(ip string) bool {
    var cmd *exec.Cmd
    
    switch runtime.GOOS {
    case "windows":
        cmd = exec.Command("ping", "-n", "1", "-w", "1000", ip)
    case "darwin":
        cmd = exec.Command("ping", "-c", "1", "-W", "1", ip)
    default: // Linux
        cmd = exec.Command("ping", "-c", "1", "-W", "1", ip)
    }
    
    err := cmd.Run()
    return err == nil
}

// getARPInfo 获取ARP信息（跨平台）
func getARPInfo(ip string) (string, string, error) {
    var mac, hostname string
    
    // 尝试获取主机名
    names, err := net.LookupAddr(ip)
    if err == nil && len(names) > 0 {
        hostname = strings.TrimSuffix(names[0], ".")
    }
    
    // 获取MAC地址（不同系统的实现）
    switch runtime.GOOS {
    case "windows":
        mac = getMACWindows(ip)
    case "linux":
        mac = getMACLinux(ip)
    case "darwin":
        mac = getMACDarwin(ip)
    default:
        // 通用方法：尝试读取系统ARP表
        mac = getMACFromARPTable(ip)
    }
    
    return mac, hostname, nil
}

// getMACFromARPTable 从系统ARP表获取MAC
func getMACFromARPTable(ip string) string {
    var cmd *exec.Cmd
    
    switch runtime.GOOS {
    case "windows":
        cmd = exec.Command("arp", "-a", ip)
    case "linux", "darwin":
        cmd = exec.Command("arp", "-n", ip)
    default:
        return ""
    }
    
    output, err := cmd.Output()
    if err != nil {
        return ""
    }
    
    outputStr := string(output)
    
    // 解析输出获取MAC地址
    lines := strings.Split(outputStr, "\n")
    for _, line := range lines {
        if strings.Contains(line, ip) {
            parts := strings.Fields(line)
            if len(parts) >= 3 {
                mac := strings.ToUpper(parts[1])
                // 验证MAC地址格式
                if isValidMAC(mac) {
                    return mac
                }
            }
        }
    }
    
    return ""
}

// 平台特定的MAC获取函数
func getMACWindows(ip string) string {
    cmd := exec.Command("arp", "-a", ip)
    output, err := cmd.Output()
    if err != nil {
        return ""
    }
    
    lines := strings.Split(string(output), "\n")
    for _, line := range lines {
        if strings.Contains(line, ip) {
            parts := strings.Fields(line)
            if len(parts) >= 2 {
                mac := strings.ToUpper(parts[1])
                if isValidMAC(mac) {
                    return mac
                }
            }
        }
    }
    
    return ""
}

func getMACLinux(ip string) string {
    cmd := exec.Command("arp", "-n", ip)
    output, err := cmd.Output()
    if err != nil {
        return ""
    }
    
    lines := strings.Split(string(output), "\n")
    for _, line := range lines {
        if strings.Contains(line, ip) {
            parts := strings.Fields(line)
            if len(parts) >= 3 {
                mac := strings.ToUpper(parts[2])
                if isValidMAC(mac) {
                    return mac
                }
            }
        }
    }
    
    return ""
}

func getMACDarwin(ip string) string {
    cmd := exec.Command("arp", "-n", ip)
    output, err := cmd.Output()
    if err != nil {
        return ""
    }
    
    lines := strings.Split(string(output), "\n")
    for _, line := range lines {
        if strings.Contains(line, ip) {
            parts := strings.Fields(line)
            if len(parts) >= 4 {
                mac := strings.ToUpper(parts[3])
                if isValidMAC(mac) {
                    return mac
                }
            }
        }
    }
    
    return ""
}

// isValidMAC 验证MAC地址格式
func isValidMAC(mac string) bool {
    // 简单的MAC地址格式验证
    return len(mac) == 17 && strings.Count(mac, ":") == 5 ||
           len(mac) == 17 && strings.Count(mac, "-") == 5
}

// getVendorFromMAC 根据MAC地址前缀获取厂商信息
func getVendorFromMAC(mac string) string {
    if len(mac) < 8 {
        return "Unknown"
    }
    
    prefix := strings.ToUpper(mac[:8])
    for vendorPrefix, vendor := range macVendors {
        if strings.HasPrefix(prefix, vendorPrefix) {
            return vendor
        }
    }
    
    return "Unknown"
}

// scanIP 扫描单个IP
func scanIP(ip string, results chan<- DeviceInfo, wg *sync.WaitGroup) {
    defer wg.Done()
    
    // Ping检测
    if !pingIP(ip) {
        return
    }
    
    // 获取MAC和主机名
    mac, hostname, err := getARPInfo(ip)
    if err != nil || mac == "" {
        return
    }
    
    // 获取厂商信息
    vendor := getVendorFromMAC(mac)
    
    // 格式化MAC地址
    mac = strings.ToUpper(mac)
    
    results <- DeviceInfo{
        IP:       ip,
        MAC:      mac,
        Hostname: hostname,
        Vendor:   vendor,
    }
}

// printTable 打印结果表格
func printTable(devices []DeviceInfo) {
    fmt.Printf("\n%-18s %-20s %-40s %s\n", "IP Address", "MAC Address", "Hostname", "Vendor")
    fmt.Println(strings.Repeat("-", 100))
    
    for _, device := range devices {
        hostname := device.Hostname
        if hostname == "" {
            hostname = "N/A"
        }
        fmt.Printf("%-18s %-20s %-40s %s\n", device.IP, device.MAC, hostname, device.Vendor)
    }
    
    fmt.Printf("\nTotal devices found: %d\n", len(devices))
}

func TestS(t *testing.T) {
    fmt.Println("=== Network Scanner ===")
    fmt.Println("Scanning local network...")
    
    // 获取本地IP和网络
    localIP, ipNet := getLocalIP()
    if localIP == nil {
        log.Fatal("Failed to get local IP address")
    }
    
    fmt.Printf("Local IP: %s\n", localIP)
    fmt.Printf("Network: %s\n", ipNet.String())
    
    // 获取要扫描的IP列表
    ips, err := getIPRange(ipNet)
    if err != nil {
        log.Fatal(err)
    }
    
    fmt.Printf("Scanning %d IP addresses...\n", len(ips))
    
    // 创建结果通道和等待组
    results := make(chan DeviceInfo, 100)
    var wg sync.WaitGroup
    
    // 启动扫描协程
    startTime := time.Now()
    maxConcurrent := 50 // 最大并发数
    semaphore := make(chan struct{}, maxConcurrent)
    
    for _, ip := range ips {
        wg.Add(1)
        semaphore <- struct{}{}
        
        go func(targetIP string) {
            defer func() { <-semaphore }()
            scanIP(targetIP, results, &wg)
        }(ip)
    }
    
    // 等待所有扫描完成
    go func() {
        wg.Wait()
        close(results)
    }()
    
    // 收集结果
    var devices []DeviceInfo
    for device := range results {
        devices = append(devices, device)
    }
    
    // 按IP排序
    sort.Slice(devices, func(i, j int) bool {
        ip1 := net.ParseIP(devices[i].IP)
        ip2 := net.ParseIP(devices[j].IP)
        return ip1.To4()[0] < ip2.To4()[0] ||
               (ip1.To4()[0] == ip2.To4()[0] && ip1.To4()[1] < ip2.To4()[1]) ||
               (ip1.To4()[0] == ip2.To4()[0] && ip1.To4()[1] == ip2.To4()[1] && ip1.To4()[2] < ip2.To4()[2]) ||
               (ip1.To4()[0] == ip2.To4()[0] && ip1.To4()[1] == ip2.To4()[1] && ip1.To4()[2] == ip2.To4()[2] && ip1.To4()[3] < ip2.To4()[3])
    })
    
    // 显示结果
    printTable(devices)
    
    duration := time.Since(startTime)
    fmt.Printf("Scan completed in %.2f seconds\n", duration.Seconds())
    
    // 保存结果到文件（可选）
    if len(devices) > 0 {
        saveToFile(devices)
    }
}

// saveToFile 保存结果到文件
func saveToFile(devices []DeviceInfo) {
    file, err := os.Create("network_devices.txt")
    if err != nil {
        log.Printf("Failed to create file: %v", err)
        return
    }
    defer file.Close()
    
    file.WriteString("IP Address,MAC Address,Hostname,Vendor\n")
    for _, device := range devices {
        hostname := device.Hostname
        if hostname == "" {
            hostname = "N/A"
        }
        file.WriteString(fmt.Sprintf("%s,%s,%s,%s\n", device.IP, device.MAC, hostname, device.Vendor))
    }
    
    fmt.Println("Results saved to network_devices.txt")
}