package engine

import (
	"bytes"
	"strings"

	"git.difuse.io/Difuse/Mellaris/analyzer"
)

type analyzerSelector struct {
	mode  AnalyzerSelectionMode
	stats *statsCounters
}

func newAnalyzerSelector(mode AnalyzerSelectionMode, stats *statsCounters) *analyzerSelector {
	if mode == "" {
		mode = AnalyzerSelectionModeSignature
	}
	return &analyzerSelector{mode: mode, stats: stats}
}

func (s *analyzerSelector) SelectTCP(ans []analyzer.Analyzer, payload []byte) []analyzer.Analyzer {
	if s == nil || s.mode == AnalyzerSelectionModeAlways || len(ans) <= 1 {
		return ans
	}
	allowed := tcpAllowedAnalyzers(payload)
	if len(allowed) == 0 {
		return ans
	}
	out := make([]analyzer.Analyzer, 0, len(ans))
	for _, a := range ans {
		name := strings.ToLower(a.Name())
		if _, known := knownTCPAnalyzers[name]; !known {
			out = append(out, a)
			continue
		}
		if allowed[name] {
			out = append(out, a)
		}
	}
	s.recordSelection(len(ans), len(out))
	if len(out) == 0 {
		return ans
	}
	return out
}

func (s *analyzerSelector) SelectUDP(ans []analyzer.Analyzer, payload []byte) []analyzer.Analyzer {
	if s == nil || s.mode == AnalyzerSelectionModeAlways || len(ans) <= 1 {
		return ans
	}
	allowed := udpAllowedAnalyzers(payload)
	if len(allowed) == 0 {
		return ans
	}
	out := make([]analyzer.Analyzer, 0, len(ans))
	for _, a := range ans {
		name := strings.ToLower(a.Name())
		if _, known := knownUDPAnalyzers[name]; !known {
			out = append(out, a)
			continue
		}
		if allowed[name] {
			out = append(out, a)
		}
	}
	s.recordSelection(len(ans), len(out))
	if len(out) == 0 {
		return ans
	}
	return out
}

func (s *analyzerSelector) recordSelection(total, selected int) {
	if s == nil || s.stats == nil || total <= 0 {
		return
	}
	s.stats.AnalyzerSelectionsTotal.Add(1)
	if selected < total {
		s.stats.AnalyzerSelectionsPruned.Add(1)
	}
}

var (
	knownTCPAnalyzers = map[string]struct{}{
		"fet":     {},
		"http":    {},
		"socks":   {},
		"ssh":     {},
		"tls":     {},
		"trojan":  {},
		"dns":     {},
		"openvpn": {},
	}
	knownUDPAnalyzers = map[string]struct{}{
		"dns":       {},
		"openvpn":   {},
		"quic":      {},
		"wireguard": {},
	}
)

func tcpAllowedAnalyzers(payload []byte) map[string]bool {
	allowed := make(map[string]bool, 4)
	if looksLikeTLS(payload) {
		allowed["tls"] = true
		allowed["trojan"] = true
		allowed["fet"] = true
	}
	if looksLikeHTTP(payload) {
		allowed["http"] = true
		allowed["fet"] = true
	}
	if looksLikeSSH(payload) {
		allowed["ssh"] = true
		allowed["fet"] = true
	}
	if looksLikeSOCKS(payload) {
		allowed["socks"] = true
		allowed["fet"] = true
	}
	if looksLikeDNSTCP(payload) {
		allowed["dns"] = true
		allowed["fet"] = true
	}
	if len(allowed) == 0 {
		return nil
	}
	return allowed
}

func udpAllowedAnalyzers(payload []byte) map[string]bool {
	allowed := make(map[string]bool, 4)
	if looksLikeWireGuard(payload) {
		allowed["wireguard"] = true
	}
	if looksLikeOpenVPN(payload) {
		allowed["openvpn"] = true
	}
	if looksLikeQUIC(payload) {
		allowed["quic"] = true
	}
	if looksLikeDNSUDP(payload) {
		allowed["dns"] = true
	}
	if len(allowed) == 0 {
		return nil
	}
	return allowed
}

func looksLikeTLS(payload []byte) bool {
	if len(payload) < 3 {
		return false
	}
	return (payload[0] == 0x16 || payload[0] == 0x17) && payload[1] == 0x03 && payload[2] <= 0x09
}

func looksLikeHTTP(payload []byte) bool {
	if len(payload) < 3 {
		return false
	}
	head := strings.ToUpper(string(payload[:3]))
	switch head {
	case "GET", "HEA", "POS", "PUT", "DEL", "CON", "OPT", "TRA", "PAT":
		return true
	default:
		return false
	}
}

func looksLikeSSH(payload []byte) bool {
	return len(payload) >= 4 && bytes.HasPrefix(payload, []byte("SSH-"))
}

func looksLikeSOCKS(payload []byte) bool {
	if len(payload) < 2 {
		return false
	}
	return payload[0] == 0x04 || payload[0] == 0x05
}

func looksLikeDNSTCP(payload []byte) bool {
	if len(payload) < 14 {
		return false
	}
	msgLen := int(payload[0])<<8 | int(payload[1])
	if msgLen <= 0 || msgLen+2 > len(payload) {
		return false
	}
	qd := int(payload[6])<<8 | int(payload[7])
	an := int(payload[8])<<8 | int(payload[9])
	return qd+an > 0
}

func looksLikeDNSUDP(payload []byte) bool {
	if len(payload) < 12 {
		return false
	}
	qd := int(payload[4])<<8 | int(payload[5])
	an := int(payload[6])<<8 | int(payload[7])
	ns := int(payload[8])<<8 | int(payload[9])
	ar := int(payload[10])<<8 | int(payload[11])
	return qd+an+ns+ar > 0
}

func looksLikeQUIC(payload []byte) bool {
	if len(payload) < 6 {
		return false
	}
	// Long header with non-zero version.
	if payload[0]&0x80 == 0 {
		return false
	}
	version := uint32(payload[1])<<24 | uint32(payload[2])<<16 | uint32(payload[3])<<8 | uint32(payload[4])
	return version != 0
}

func looksLikeOpenVPN(payload []byte) bool {
	if len(payload) == 0 {
		return false
	}
	opcode := payload[0] >> 3
	return opcode >= 1 && opcode <= 11
}

func looksLikeWireGuard(payload []byte) bool {
	if len(payload) < 4 {
		return false
	}
	if payload[0] < 1 || payload[0] > 4 {
		return false
	}
	return payload[1] == 0 && payload[2] == 0 && payload[3] == 0
}