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)
				}()

				// 检查发包限制
				if canSend, reason := common.CanSendPacket(); !canSend {
					common.LogError(fmt.Sprintf("端口扫描 %s 受限: %s", addr, reason))
					return nil
				}
				
				// 连接测试 - 支持SOCKS5代理
				conn, err := common.WrapperTcpWithTimeout("tcp", addr, to)
				
				if err != nil {
					// 计数TCP连接失败包
					common.IncrementTCPFailedPacketCount()
					return nil
				}
				defer conn.Close()
				
				// 计数TCP连接成功包
				common.IncrementTCPSuccessPacketCount()

				// 记录开放端口（先记录到内存，延迟输出直到服务识别完成）
				atomic.AddInt64(&count, 1)
				aliveMap.Store(addr, struct{}{})
				common.SaveResult(&output.ScanResult{
					Time: time.Now(), Type: output.TypePort, Target: host,
					Status: "open", Details: map[string]interface{}{"port": port},
				})

				// 智能服务识别：保持准确性，优化网络交互
				serviceInfo, err := NewSmartPortInfoScanner(host, port, conn, to).SmartIdentify()
				if err == nil {
					// 构建结果详情
					details := map[string]interface{}{"port": port, "service": serviceInfo.Name}
					if serviceInfo.Version != "" {
						details["version"] = serviceInfo.Version
					}

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

					// 智能判断是否为Web服务
					isWeb := IsWebServiceByFingerprint(serviceInfo)
					if isWeb {
						details["is_web"] = true
						// 标记该端口为Web服务，后续会自动启用WebTitle和WebPOC
						MarkAsWebService(host, port, serviceInfo)
					}

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

					// 构建统一的服务信息日志
					var sb strings.Builder
					sb.WriteString("端口开放 " + addr)
					
					if serviceInfo.Name != "unknown" {
						sb.WriteString(" [" + serviceInfo.Name + "]")
						if isWeb {
							sb.WriteString("(Web服务)")
						}
					}
					
					if serviceInfo.Version != "" {
						sb.WriteString(" 版本:" + serviceInfo.Version)
					}

					// 添加详细服务信息（默认启用）
					for k, v := range serviceInfo.Extras {
						if v == "" {
							continue
						}
						switch k {
						case "vendor_product":
							sb.WriteString(" 产品:" + v)
						case "os":
							sb.WriteString(" 系统:" + v)
						case "info":
							sb.WriteString(" 信息:" + v)
						}
					}

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

					// 统一输出端口和服务信息
					common.LogInfo(sb.String())
				} else {
					// 服务识别失败时，只输出端口开放信息
					common.LogInfo("端口开放 " + addr)
					common.LogDebug(fmt.Sprintf("服务识别失败 %s: %v，回退到Web检测", addr, err))
				}

				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
}