package core

import (
	"fmt"
	"github.com/shadow1ng/fscan/common"
	"net"
	"time"
)

// SmartPortInfoScanner 智能服务识别器：保持nmap准确性，优化网络交互
type SmartPortInfoScanner struct {
	Address string
	Port    int
	Conn    net.Conn
	Timeout time.Duration
	info    *Info
}

// NewSmartPortInfoScanner 创建智能服务识别器
func NewSmartPortInfoScanner(addr string, port int, conn net.Conn, timeout time.Duration) *SmartPortInfoScanner {
	return &SmartPortInfoScanner{
		Address: addr,
		Port:    port,
		Conn:    conn,
		Timeout: timeout,
		info: &Info{
			Address: addr,
			Port:    port,
			Conn:    conn,
			Result: Result{
				Service: Service{},
			},
		},
	}
}

// SmartIdentify 智能服务识别：Banner优先 + 优化的探测策略
func (s *SmartPortInfoScanner) SmartIdentify() (*ServiceInfo, error) {
	common.LogDebug(fmt.Sprintf("开始智能识别服务 %s:%d", s.Address, s.Port))
	
	// 第一阶段：读取初始Banner（大部分服务会主动发送）
	s.tryInitialBanner()
	if s.info.Found {
		return s.buildServiceInfo(), nil
	}
	
	// 第二阶段：智能探测策略（减少探测器数量）
	s.smartProbeStrategy()
	
	// 构造返回结果
	return s.buildServiceInfo(), nil
}

// tryInitialBanner 尝试读取服务主动发送的Banner
func (s *SmartPortInfoScanner) tryInitialBanner() {
	// 读取初始响应
	if response, err := s.info.Read(); err == nil && len(response) > 0 {
		common.LogDebug(fmt.Sprintf("收到初始Banner: %d 字节", len(response)))
		
		// 使用原有的nmap指纹库解析Banner，保持准确性
		if s.info.tryProbes(response, []*Probe{null, commonProbe}) {
			common.LogDebug("Banner识别成功")
			return
		}
	}
}

// smartProbeStrategy 智能探测策略：减少探测器数量但保持准确性
func (s *SmartPortInfoScanner) smartProbeStrategy() {
	// 记录已使用的探测器
	usedProbes := make(map[string]struct{})
	
	// 优先使用端口专用的第一个探测器（而不是全部）
	if s.tryFirstPortProbe(usedProbes) {
		return
	}
	
	// 如果端口专用探测失败，使用最有效的通用探测器
	s.tryBestGenericProbes(usedProbes)
}

// tryFirstPortProbe 只尝试该端口的第一个专用探测器
func (s *SmartPortInfoScanner) tryFirstPortProbe(usedProbes map[string]struct{}) bool {
	// 检查是否有端口专用探测器
	portProbes := common.PortMap[s.info.Port]
	if len(portProbes) == 0 {
		return false
	}
	
	// 只使用第一个（最可能的）探测器
	probeName := portProbes[0]
	usedProbes[probeName] = struct{}{}
	
	probe := v.ProbesMapKName[probeName]
	probeData, err := DecodeData(probe.Data)
	if err != nil || len(probeData) == 0 {
		return false
	}
	
	common.LogDebug(fmt.Sprintf("使用首选探测器: %s", probeName))
	if response := s.info.Connect(probeData); len(response) > 0 {
		// 使用当前探测器检查响应
		s.info.GetInfo(response, &probe)
		if s.info.Found {
			return true
		}
		
		// 也用通用探测器检查这个响应
		if probeName != "NULL" {
			return s.info.tryProbes(response, []*Probe{commonProbe})
		}
	}
	
	return false
}

// tryBestGenericProbes 使用最有效的通用探测器（限制数量）
func (s *SmartPortInfoScanner) tryBestGenericProbes(usedProbes map[string]struct{}) {
	// 选择最有效的探测器：GetRequest适用于HTTP，GenericLines适用于大多数文本协议
	bestProbes := []string{"GetRequest", "GenericLines", "HTTPOptions"}
	
	for _, probeName := range bestProbes {
		// 跳过已使用的探测器
		if _, used := usedProbes[probeName]; used {
			continue
		}
		
		probe, exists := v.ProbesMapKName[probeName]
		if !exists {
			continue
		}
		
		probeData, err := DecodeData(probe.Data)
		if err != nil || len(probeData) == 0 {
			continue
		}
		
		common.LogDebug(fmt.Sprintf("使用通用探测器: %s", probeName))
		response := s.info.Connect(probeData)
		if len(response) == 0 {
			continue
		}
		
		// 检查当前探测器
		s.info.GetInfo(response, &probe)
		if s.info.Found {
			return
		}
		
		// 也检查通用匹配
		if probeName != "GenericLines" {
			if s.info.tryProbes(response, []*Probe{commonProbe}) {
				return
			}
		}
	}
	
	// 如果所有探测都失败，标记为未知服务
	if s.info.Result.Service.Name == "" {
		s.info.Result.Service.Name = "unknown"
	}
}

// buildServiceInfo 构建ServiceInfo结果
func (s *SmartPortInfoScanner) buildServiceInfo() *ServiceInfo {
	result := &s.info.Result
	
	serviceInfo := &ServiceInfo{
		Name:    result.Service.Name,
		Banner:  result.Banner,
		Version: result.Service.Extras["version"],
		Extras:  make(map[string]string),
	}
	
	// 复制额外信息
	for k, v := range result.Service.Extras {
		serviceInfo.Extras[k] = v
	}
	
	common.LogDebug(fmt.Sprintf("智能识别完成 %s:%d => %s", s.Address, s.Port, serviceInfo.Name))
	return serviceInfo
}