package net

import (
	"b612.me/apps/b612/netforward"
	"b612.me/stario"
	"b612.me/starlog"
	"b612.me/starnet"
	"fmt"
	"math"
	"net"
	"sync/atomic"
	"time"
)

type ScanIP struct {
	Host         string
	CIDR         int
	Port         int
	Mask         string
	Threads      int
	Timeout      int
	ScanType     string
	ipNet        *net.IPNet
	Log          string
	Retry        int
	WithHostname bool
}

func (s *ScanIP) Parse() error {
	if s.CIDR == 0 && s.Mask == "" {
		return fmt.Errorf("CIDR or Mask must be set")

	}
	if s.CIDR != 0 {
		return nil
	}
	//mask to cidr
	ipMask := net.IPMask(net.ParseIP(s.Mask).To4())
	if ipMask == nil {
		return fmt.Errorf("invalid mask")
	}
	s.CIDR, _ = ipMask.Size()
	return nil
}

func (s *ScanIP) nextIP(ipStr string) (net.IP, error) {
	var err error
	if s.ipNet == nil {
		_, s.ipNet, err = net.ParseCIDR(s.Host + "/" + fmt.Sprintf("%d", s.CIDR))
		if err != nil {
			return nil, fmt.Errorf("invalid CIDR: %v", err)
		}
	}
	ip := net.ParseIP(ipStr)
	if ip == nil {
		return nil, fmt.Errorf("invalid IP: %v", ipStr)
	}

	// Convert IP to 4-byte representation
	ip = ip.To4()
	if ip == nil {
		return nil, fmt.Errorf("non-IPv4 address: %v", ipStr)
	}

	// Increment IP
	for i := len(ip) - 1; i >= 0; i-- {
		ip[i]++
		if ip[i] > 0 {
			break
		}
	}

	// Check if incremented IP is still in range
	if !s.ipNet.Contains(ip) {
		return nil, nil
	}

	return ip, nil
}

func (s *ScanIP) NetSize() (int, error) {
	var err error
	if s.ipNet == nil {
		_, s.ipNet, err = net.ParseCIDR(s.Host + "/" + fmt.Sprintf("%d", s.CIDR))
		if err != nil {
			return 0, fmt.Errorf("invalid CIDR: %v", err)
		}
	}

	maskSize, _ := s.ipNet.Mask.Size()
	return int(math.Pow(2, float64(32-maskSize))) - 2, nil
}

func (s *ScanIP) FirstLastIP() (net.IP, net.IP, error) {
	var err error
	if s.ipNet == nil {
		_, s.ipNet, err = net.ParseCIDR(s.Host + "/" + fmt.Sprintf("%d", s.CIDR))
		if err != nil {
			return nil, nil, fmt.Errorf("invalid CIDR: %v", err)
		}
	}

	firstIP := s.ipNet.IP.Mask(s.ipNet.Mask)
	lastIP := make(net.IP, len(firstIP))
	copy(lastIP, firstIP)
	for i := range firstIP {
		lastIP[i] = firstIP[i] | ^s.ipNet.Mask[i]
	}

	return firstIP, lastIP, nil
}

func (s *ScanIP) ICMP() error {
	if s.ScanType != "icmp" {
		return fmt.Errorf("scan type must be icmp")
	}
	if err := s.Parse(); err != nil {
		return err
	}
	if s.Log != "" {
		starlog.SetLogFile(s.Log, starlog.Std, true)
	}
	firstIP, lastIP, err := s.FirstLastIP()
	if err != nil {
		return err
	}
	ns, _ := s.NetSize()
	starlog.Infof("Scan %s/%d\n", s.Host, s.CIDR)
	starlog.Infof("Scan %s-%s\n", firstIP.String(), lastIP.String())
	starlog.Infof("There are %d hosts\n", ns)
	starlog.Infof("Threads: %d\n", s.Threads)

	wg := stario.NewWaitGroup(s.Threads)
	count := int32(0)
	allcount := int32(0)
	interrupt := make(chan [2]string)
	go func() {
		for {
			select {
			case <-time.After(time.Second * 2):
				fmt.Printf("scan %d ips, %d up\r", allcount, count)
			case ip, opened := <-interrupt:
				if !opened {
					return
				}
				if s.WithHostname {
					starlog.Infof("Host %v is up, Name:%v\n", ip[0], ip[1])
				} else {
					starlog.Infof("Host %v is up\n", ip[0])
				}
			}

		}
	}()
	idx := 0
	for {
		ip := firstIP.String()
		if ip == lastIP.String() {
			break
		}
		idx++
		wg.Add(1)
		go func(ip string, idx int) {
			defer func() {
				atomic.AddInt32(&allcount, 1)
			}()
			defer wg.Done()
			for i := 0; i < s.Retry+1; i++ {
				_, err := starnet.Ping(ip, idx, time.Duration(s.Timeout)*time.Millisecond)
				if err == nil {
					atomic.AddInt32(&count, 1)
					if s.WithHostname {
						hostname, err := net.LookupAddr(ip)
						if err == nil {
							interrupt <- [2]string{ip, hostname[0]}
							return
						}
					}
					interrupt <- [2]string{ip, ""}
					return
				}
			}
		}(ip, idx)
		firstIP, _ = s.nextIP(ip)
	}
	wg.Wait()
	close(interrupt)
	starlog.Infof("scan %d ips, %d up\n", ns, count)
	return nil
}

func (s *ScanIP) TCP(port int) error {
	if s.ScanType != "tcp" {
		return fmt.Errorf("scan type must be tcp")
	}
	if err := s.Parse(); err != nil {
		return err
	}
	if s.Log != "" {
		starlog.SetLogFile(s.Log, starlog.Std, true)
	}
	firstIP, lastIP, err := s.FirstLastIP()
	if err != nil {
		return err
	}
	ns, _ := s.NetSize()
	starlog.Infof("Scan %s/%d\n", s.Host, s.CIDR)
	starlog.Infof("Scan %s-%s\n", firstIP.String(), lastIP.String())
	starlog.Infof("There are %d hosts\n", ns)
	starlog.Infof("Threads: %d\n", s.Threads)

	wg := stario.NewWaitGroup(s.Threads)
	count := int32(0)
	allcount := int32(0)
	interrupt := make(chan [2]string)
	go func() {
		for {
			select {
			case <-time.After(time.Second * 2):
				fmt.Printf("scan %d ips, %d up\r", allcount, count)
			case ip, opened := <-interrupt:
				if !opened {
					return
				}
				if s.WithHostname {
					starlog.Infof("Host %v is up, Name:%v\n", ip[0], ip[1])
				} else {
					starlog.Infof("Host %v is up\n", ip[0])
				}
			}

		}
	}()
	idx := 0
	localAddr, err := net.ResolveTCPAddr("tcp", ":0")
	if err != nil {
		starlog.Errorln("ResolveTCPAddr error, ", err)
		return err
	}
	for {
		ip := firstIP.String()
		if ip == lastIP.String() {
			break
		}
		idx++
		wg.Add(1)
		go func(ip string, idx int) {
			defer func() {
				atomic.AddInt32(&allcount, 1)
			}()
			defer wg.Done()
			for i := 0; i < s.Retry+1; i++ {
				dialer := net.Dialer{
					LocalAddr: localAddr,
					Timeout:   time.Duration(s.Timeout) * time.Millisecond,
					Control:   netforward.ControlSetReUseAddr,
				}
				conn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", ip, port))
				if err == nil {
					conn.(*net.TCPConn).SetLinger(0)
					conn.Close()
					atomic.AddInt32(&count, 1)
					if s.WithHostname {
						hostname, err := net.LookupAddr(ip)
						if err == nil {
							interrupt <- [2]string{ip, hostname[0]}
							return
						}
					}
					interrupt <- [2]string{ip, ""}
					return
				}
			}
		}(ip, idx)
		firstIP, _ = s.nextIP(ip)
	}
	wg.Wait()
	close(interrupt)
	starlog.Infof("scan %d ips, %d up\n", ns, count)
	return nil
}