package config

import (
	"sort"
	"sync"
)

// ProbeMapping 探测器映射管理器
type ProbeMapping struct {
	mu          sync.RWMutex
	defaultMap  []string
	portMap     map[int][]string
	initialized bool
}

// NewProbeMapping 创建探测器映射管理器
func NewProbeMapping() *ProbeMapping {
	return &ProbeMapping{
		defaultMap:  getDefaultProbeMap(),
		portMap:     getDefaultPortMap(),
		initialized: true,
	}
}

// getDefaultProbeMap 获取默认的探测器顺序
func getDefaultProbeMap() []string {
	return []string{
		"GenericLines",
		"GetRequest",
		"TLSSessionReq",
		"SSLSessionReq",
		"ms-sql-s",
		"JavaRMI",
		"LDAPSearchReq",
		"LDAPBindReq",
		"oracle-tns",
		"Socks5",
	}
}

// getDefaultPortMap 获取默认的端口映射
func getDefaultPortMap() map[int][]string {
	return map[int][]string{
		1:     {"GetRequest", "Help"},
		7:     {"Help"},
		21:    {"GenericLines", "Help"},
		23:    {"GenericLines", "tn3270"},
		25:    {"Hello", "Help"},
		35:    {"GenericLines"},
		42:    {"SMBProgNeg"},
		43:    {"GenericLines"},
		53:    {"DNSVersionBindReqTCP", "DNSStatusRequestTCP"},
		70:    {"GetRequest"},
		79:    {"GenericLines", "GetRequest", "Help"},
		80:    {"GetRequest", "HTTPOptions", "RTSPRequest", "X11Probe", "FourOhFourRequest"},
		81:    {"GetRequest", "HTTPOptions", "RPCCheck", "FourOhFourRequest"},
		82:    {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		83:    {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		84:    {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		85:    {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		88:    {"GetRequest", "Kerberos", "SMBProgNeg", "FourOhFourRequest"},
		98:    {"GenericLines"},
		110:   {"GenericLines"},
		111:   {"RPCCheck"},
		113:   {"GenericLines", "GetRequest", "Help"},
		119:   {"GenericLines", "Help"},
		130:   {"NotesRPC"},
		135:   {"DNSVersionBindReqTCP", "SMBProgNeg"},
		139:   {"GetRequest", "SMBProgNeg"},
		143:   {"GetRequest"},
		175:   {"NJE"},
		199:   {"GenericLines", "RPCCheck", "Socks5", "Socks4"},
		214:   {"GenericLines"},
		264:   {"GenericLines"},
		311:   {"LDAPSearchReq"},
		340:   {"GenericLines"},
		389:   {"LDAPSearchReq", "LDAPBindReq"},
		443:   {"GetRequest", "HTTPOptions", "SSLSessionReq", "TerminalServerCookie", "TLSSessionReq"},
		444:   {"GetRequest", "HTTPOptions", "SSLSessionReq", "TerminalServerCookie", "TLSSessionReq"},
		445:   {"SMBProgNeg"},
		465:   {"Hello", "Help", "GetRequest", "HTTPOptions", "SSLSessionReq", "TerminalServerCookie"},
		502:   {"GenericLines"},
		503:   {"GenericLines"},
		513:   {"GenericLines"},
		514:   {"GenericLines"},
		515:   {"LPDString"},
		544:   {"GenericLines"},
		548:   {"afp"},
		554:   {"GetRequest"},
		563:   {"GenericLines"},
		587:   {"Hello", "Help"},
		631:   {"GetRequest", "HTTPOptions"},
		636:   {"LDAPSearchReq", "LDAPBindReq", "SSLSessionReq"},
		646:   {"LDAPSearchReq", "RPCCheck"},
		691:   {"GenericLines"},
		873:   {"GenericLines"},
		898:   {"GetRequest"},
		993:   {"GenericLines", "SSLSessionReq", "TerminalServerCookie", "TLSSessionReq"},
		995:   {"GenericLines", "SSLSessionReq", "TerminalServerCookie", "TLSSessionReq"},
		1080:  {"GenericLines", "Socks5", "Socks4"},
		1099:  {"JavaRMI"},
		1234:  {"SqueezeCenter_CLI"},
		1311:  {"GenericLines"},
		1352:  {"oracle-tns"},
		1414:  {"ibm-mqseries"},
		1433:  {"ms-sql-s"},
		1521:  {"oracle-tns"},
		1723:  {"GenericLines"},
		1883:  {"mqtt"},
		1911:  {"oracle-tns"},
		2000:  {"GenericLines", "oracle-tns"},
		2049:  {"RPCCheck"},
		2121:  {"GenericLines", "Help"},
		2181:  {"GenericLines"},
		2222:  {"GetRequest", "GenericLines", "HTTPOptions", "Help", "SSH", "TerminalServerCookie"},
		2375:  {"docker", "GetRequest", "HTTPOptions"},
		2376:  {"docker", "GetRequest", "HTTPOptions", "SSLSessionReq"},
		2484:  {"oracle-tns"},
		2628:  {"dominoconsole"},
		3000:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		3268:  {"LDAPSearchReq", "LDAPBindReq"},
		3269:  {"LDAPSearchReq", "LDAPBindReq", "SSLSessionReq"},
		3306:  {"GenericLines", "GetRequest", "HTTPOptions"},
		3389:  {"TerminalServerCookie", "TerminalServer"},
		3690:  {"GenericLines"},
		4000:  {"GenericLines"},
		4369:  {"epmd"},
		4444:  {"GenericLines"},
		4840:  {"GenericLines"},
		5000:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		5050:  {"GenericLines"},
		5060:  {"SIPOptions"},
		5222:  {"GenericLines"},
		5432:  {"GenericLines"},
		5555:  {"GenericLines"},
		5560:  {"GenericLines", "oracle-tns"},
		5631:  {"GenericLines", "PCWorkstation"},
		5672:  {"GenericLines"},
		5984:  {"GetRequest", "HTTPOptions"},
		6000:  {"X11Probe"},
		6379:  {"redis-server"},
		6432:  {"GenericLines"},
		6667:  {"GenericLines"},
		7000:  {"GetRequest", "HTTPOptions", "FourOhFourRequest", "JavaRMI"},
		7001:  {"GetRequest", "HTTPOptions", "FourOhFourRequest", "JavaRMI"},
		7002:  {"GetRequest", "HTTPOptions", "FourOhFourRequest", "JavaRMI"},
		7070:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		7443:  {"GetRequest", "HTTPOptions", "SSLSessionReq"},
		7777:  {"GenericLines", "oracle-tns"},
		8000:  {"GetRequest", "HTTPOptions", "FourOhFourRequest", "iperf3"},
		8005:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		8008:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		8009:  {"GetRequest", "HTTPOptions", "FourOhFourRequest", "ajp"},
		8080:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		8081:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		8089:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		8090:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		8443:  {"GetRequest", "HTTPOptions", "SSLSessionReq"},
		8888:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		9000:  {"GetRequest", "HTTPOptions", "FourOhFourRequest"},
		9042:  {"GenericLines"},
		9092:  {"GenericLines", "kafka"},
		9200:  {"GetRequest", "HTTPOptions", "elasticsearch"},
		9300:  {"GenericLines"},
		9999:  {"GetRequest", "HTTPOptions", "FourOhFourRequest", "adbConnect"},
		10000: {"GetRequest", "HTTPOptions", "FourOhFourRequest", "JavaRMI"},
		10051: {"GenericLines"},
		11211: {"Memcache"},
		15672: {"GetRequest", "HTTPOptions"},
		27017: {"mongodb"},
		27018: {"mongodb"},
		50070: {"GetRequest", "HTTPOptions"},
		61616: {"GenericLines"},
	}
}

// GetProbesForPort 获取指定端口的探测器列表
func (pm *ProbeMapping) GetProbesForPort(port int) []string {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	if probes, exists := pm.portMap[port]; exists {
		// 返回副本，避免外部修改
		result := make([]string, len(probes))
		copy(result, probes)
		return result
	}
	
	// 返回默认探测器
	result := make([]string, len(pm.defaultMap))
	copy(result, pm.defaultMap)
	return result
}

// SetProbesForPort 设置指定端口的探测器列表
func (pm *ProbeMapping) SetProbesForPort(port int, probes []string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	if len(probes) > 0 {
		// 创建副本
		probesCopy := make([]string, len(probes))
		copy(probesCopy, probes)
		pm.portMap[port] = probesCopy
	} else {
		delete(pm.portMap, port)
	}
}

// AddProbeToPort 向指定端口添加探测器
func (pm *ProbeMapping) AddProbeToPort(port int, probe string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	if probe == "" {
		return
	}
	
	if existing, exists := pm.portMap[port]; exists {
		// 检查是否已存在
		for _, p := range existing {
			if p == probe {
				return
			}
		}
		// 添加到现有列表
		pm.portMap[port] = append(existing, probe)
	} else {
		// 创建新列表
		pm.portMap[port] = []string{probe}
	}
}

// RemoveProbeFromPort 从指定端口移除探测器
func (pm *ProbeMapping) RemoveProbeFromPort(port int, probe string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	if existing, exists := pm.portMap[port]; exists {
		var result []string
		for _, p := range existing {
			if p != probe {
				result = append(result, p)
			}
		}
		
		if len(result) > 0 {
			pm.portMap[port] = result
		} else {
			delete(pm.portMap, port)
		}
	}
}

// GetDefaultProbes 获取默认探测器列表
func (pm *ProbeMapping) GetDefaultProbes() []string {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	result := make([]string, len(pm.defaultMap))
	copy(result, pm.defaultMap)
	return result
}

// SetDefaultProbes 设置默认探测器列表
func (pm *ProbeMapping) SetDefaultProbes(probes []string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	if len(probes) > 0 {
		defaultCopy := make([]string, len(probes))
		copy(defaultCopy, probes)
		pm.defaultMap = defaultCopy
	}
}

// GetAllPortMappings 获取所有端口映射
func (pm *ProbeMapping) GetAllPortMappings() map[int][]string {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	result := make(map[int][]string)
	for port, probes := range pm.portMap {
		probesCopy := make([]string, len(probes))
		copy(probesCopy, probes)
		result[port] = probesCopy
	}
	return result
}

// GetMappedPorts 获取所有已映射的端口列表
func (pm *ProbeMapping) GetMappedPorts() []int {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	ports := make([]int, 0, len(pm.portMap))
	for port := range pm.portMap {
		ports = append(ports, port)
	}
	
	sort.Ints(ports)
	return ports
}

// HasMappingForPort 检查指定端口是否有自定义映射
func (pm *ProbeMapping) HasMappingForPort(port int) bool {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	_, exists := pm.portMap[port]
	return exists
}

// GetProbeCount 获取指定端口的探测器数量
func (pm *ProbeMapping) GetProbeCount(port int) int {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	if probes, exists := pm.portMap[port]; exists {
		return len(probes)
	}
	return len(pm.defaultMap)
}

// ResetToDefaults 重置为默认映射
func (pm *ProbeMapping) ResetToDefaults() {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	pm.defaultMap = getDefaultProbeMap()
	pm.portMap = getDefaultPortMap()
}

// ClearPortMapping 清除指定端口的映射
func (pm *ProbeMapping) ClearPortMapping(port int) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	delete(pm.portMap, port)
}

// ClearAllMappings 清除所有端口映射
func (pm *ProbeMapping) ClearAllMappings() {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	pm.portMap = make(map[int][]string)
}

// 全局探测器映射实例
var (
	globalProbeMapping *ProbeMapping
	probeMappingOnce   sync.Once
)

// GetGlobalProbeMapping 获取全局探测器映射实例
func GetGlobalProbeMapping() *ProbeMapping {
	probeMappingOnce.Do(func() {
		globalProbeMapping = NewProbeMapping()
	})
	return globalProbeMapping
}

// SetGlobalProbeMapping 设置全局探测器映射实例
func SetGlobalProbeMapping(mapping *ProbeMapping) {
	globalProbeMapping = mapping
}

// 便利函数，直接使用全局实例

// GetProbesForPort 获取指定端口的探测器列表（全局函数）
func GetProbesForPort(port int) []string {
	return GetGlobalProbeMapping().GetProbesForPort(port)
}

// GetDefaultProbes 获取默认探测器列表（全局函数）
func GetDefaultProbes() []string {
	return GetGlobalProbeMapping().GetDefaultProbes()
}

// SetProbesForPort 设置指定端口的探测器列表（全局函数）
func SetProbesForPort(port int, probes []string) {
	GetGlobalProbeMapping().SetProbesForPort(port, probes)
}

// GetMappedPorts 获取所有已映射的端口列表（全局函数）
func GetMappedPorts() []int {
	return GetGlobalProbeMapping().GetMappedPorts()
}