package core

import (
	"fmt"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/common/i18n"
	"github.com/shadow1ng/fscan/common/parsers"
	"strings"
)

/*
PortDiscoveryService.go - 端口发现服务

负责主机存活检测、端口扫描等端口发现相关功能，
从ServiceScanner中分离出来提高代码可维护性。
*/

// PortDiscoveryService 端口发现服务
type PortDiscoveryService struct{}

// NewPortDiscoveryService 创建端口发现服务
func NewPortDiscoveryService() *PortDiscoveryService {
	return &PortDiscoveryService{}
}

// DiscoverTargets 发现目标主机和端口
func (p *PortDiscoveryService) DiscoverTargets(hostInput string, baseInfo common.HostInfo) ([]common.HostInfo, error) {
	// 解析目标主机
	hosts, err := parsers.ParseIP(hostInput, common.HostsFile, common.ExcludeHosts)
	if err != nil {
		return nil, fmt.Errorf(i18n.GetText("parse_error_target_failed"), err)
	}

	var targetInfos []common.HostInfo

	// 主机存活性检测和端口扫描
	if len(hosts) > 0 || len(common.HostPort) > 0 {
		// 主机存活检测
		if p.shouldPerformLivenessCheck(hosts) {
			hosts = CheckLive(hosts, false)
			common.LogBase(i18n.GetText("scan_alive_hosts_count", len(hosts)))
		}

		// 端口扫描
		alivePorts := p.discoverAlivePorts(hosts)
		if len(alivePorts) > 0 {
			targetInfos = p.convertToTargetInfos(alivePorts, baseInfo)
		}
	}

	return targetInfos, nil
}

// shouldPerformLivenessCheck 判断是否需要执行存活性检测
func (p *PortDiscoveryService) shouldPerformLivenessCheck(hosts []string) bool {
	return common.DisablePing == false && len(hosts) > 1
}

// discoverAlivePorts 发现存活的端口
func (p *PortDiscoveryService) discoverAlivePorts(hosts []string) []string {
	var alivePorts []string

	// 根据扫描模式选择端口扫描方式
	if len(hosts) > 0 {
		alivePorts = EnhancedPortScan(hosts, common.Ports, common.Timeout)
		common.LogBase(i18n.GetText("scan_alive_ports_count", len(alivePorts)))
	}

	// UDP端口特殊处理（当前仅支持SNMP的161端口）
	udpPorts := p.handleUDPPorts(hosts)
	if len(udpPorts) > 0 {
		alivePorts = append(alivePorts, udpPorts...)
		common.LogBase(i18n.GetText("scan_alive_ports_count", len(alivePorts)))
	}

	// 合并额外指定的端口
	if len(common.HostPort) > 0 {
		alivePorts = append(alivePorts, common.HostPort...)
		alivePorts = common.RemoveDuplicate(alivePorts)
		common.HostPort = nil
		common.LogBase(i18n.GetText("scan_alive_ports_count", len(alivePorts)))
	}

	return alivePorts
}

// convertToTargetInfos 将端口列表转换为目标信息
func (p *PortDiscoveryService) convertToTargetInfos(ports []string, baseInfo common.HostInfo) []common.HostInfo {
	var infos []common.HostInfo

	for _, targetIP := range ports {
		hostParts := strings.Split(targetIP, ":")
		if len(hostParts) != 2 {
			common.LogError(i18n.GetText("parse_error_invalid_target_format", targetIP))
			continue
		}

		info := baseInfo
		info.Host = hostParts[0]
		info.Ports = hostParts[1]
		infos = append(infos, info)
	}

	return infos
}

// handleUDPPorts 处理UDP端口的特殊逻辑
func (p *PortDiscoveryService) handleUDPPorts(hosts []string) []string {
	var udpPorts []string
	
	// 检查是否包含SNMP端口161
	portList := parsers.ParsePort(common.Ports)
	hasPort161 := false
	for _, port := range portList {
		if port == 161 {
			hasPort161 = true
			break
		}
	}
	
	// 如果端口列表包含161，则为每个主机添加UDP 161端口
	if hasPort161 {
		for _, host := range hosts {
			udpPorts = append(udpPorts, fmt.Sprintf("%s:161", host))
		}
		if len(udpPorts) > 0 {
			common.LogBase(i18n.GetText("scan_snmp_udp_ports_added"))
		}
	}
	
	return udpPorts
}