package core

import (
	"context"
	"fmt"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/common/i18n"
	"github.com/shadow1ng/fscan/common/output"
	"github.com/shadow1ng/fscan/common/parsers"
	"golang.org/x/sync/errgroup"
	"golang.org/x/sync/semaphore"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// EnhancedPortScan 高性能端口扫描函数
func EnhancedPortScan(hosts []string, ports string, timeout int64) []string {
	// 解析端口和排除端口
	portList := parsers.ParsePort(ports)
	if len(portList) == 0 {
		common.LogError("无效端口: " + ports)
		return nil
	}

	// 预估排除端口数量，通常不会超过100个
	excludePorts := parsers.ParsePort(common.ExcludePorts)
	exclude := make(map[int]struct{}, len(excludePorts))
	for _, p := range excludePorts {
		exclude[p] = struct{}{}
	}

	// 计算总扫描数量
	totalTasks := 0
	for range hosts {
		for _, port := range portList {
			if _, excluded := exclude[port]; !excluded {
				totalTasks++
			}
		}
	}

	// 初始化端口扫描进度条
	if totalTasks > 0 && common.ShowProgress {
		description := i18n.GetText("progress_port_scanning_with_threads", common.ThreadNum)
		common.InitProgressBar(int64(totalTasks), description)
	}

	// 初始化并发控制
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	to := time.Duration(timeout) * time.Second
	sem := semaphore.NewWeighted(int64(common.ThreadNum))
	var count int64
	var aliveMap sync.Map
	g, ctx := errgroup.WithContext(ctx)

	// 并发扫描所有目标
	for _, host := range hosts {
		for _, port := range portList {
			if _, excluded := exclude[port]; excluded {
				continue
			}

			host, port := host, port // 捕获循环变量
			addr := fmt.Sprintf("%s:%d", host, port)

			if err := sem.Acquire(ctx, 1); err != nil {
				break
			}

			g.Go(func() error {
				defer func() {
					sem.Release(1)
					// 更新端口扫描进度
					common.UpdateProgressBar(1)
				}()

				// 连接测试 - 支持SOCKS5代理
				conn, err := common.WrapperTcpWithTimeout("tcp", addr, to)
				if err != nil {
					return nil
				}
				defer conn.Close()

				// 记录开放端口
				atomic.AddInt64(&count, 1)
				aliveMap.Store(addr, struct{}{})
				common.LogInfo("端口开放 " + addr)
				common.SaveResult(&output.ScanResult{
					Time: time.Now(), Type: output.TypePort, Target: host,
					Status: "open", Details: map[string]interface{}{"port": port},
				})

				// 服务识别
				if common.EnableFingerprint {
					if info, err := NewPortInfoScanner(host, port, conn, to).Identify(); err == nil {
						// 构建结果详情
						details := map[string]interface{}{"port": port, "service": info.Name}
						if info.Version != "" {
							details["version"] = info.Version
						}

						// 处理额外信息
						for k, v := range info.Extras {
							if v == "" {
								continue
							}
							switch k {
							case "vendor_product":
								details["product"] = v
							case "os", "info":
								details[k] = v
							}
						}
						if len(info.Banner) > 0 {
							details["banner"] = strings.TrimSpace(info.Banner)
						}

						// 保存服务结果
						common.SaveResult(&output.ScanResult{
							Time: time.Now(), Type: output.TypeService, Target: host,
							Status: "identified", Details: details,
						})

						// 记录服务信息
						var sb strings.Builder
						sb.WriteString("服务识别 " + addr + " => ")
						if info.Name != "unknown" {
							sb.WriteString("[" + info.Name + "]")
						}
						if info.Version != "" {
							sb.WriteString(" 版本:" + info.Version)
						}

						for k, v := range info.Extras {
							if v == "" {
								continue
							}
							switch k {
							case "vendor_product":
								sb.WriteString(" 产品:" + v)
							case "os":
								sb.WriteString(" 系统:" + v)
							case "info":
								sb.WriteString(" 信息:" + v)
							}
						}

						if len(info.Banner) > 0 && len(info.Banner) < 100 {
							sb.WriteString(" Banner:[" + strings.TrimSpace(info.Banner) + "]")
						}

						common.LogInfo(sb.String())
					}
				}

				return nil
			})
		}
	}

	_ = g.Wait()

	// 收集结果
	var aliveAddrs []string
	aliveMap.Range(func(key, _ interface{}) bool {
		aliveAddrs = append(aliveAddrs, key.(string))
		return true
	})

	// 完成端口扫描进度条
	if common.IsProgressActive() {
		common.FinishProgressBar()
	}
	
	common.LogBase(i18n.GetText("scan_complete_ports_found", count))
	return aliveAddrs
}