sub-store/internal/util/flag.go

package util

import (
	"net"
	"regexp"
	"strings"

	"github.com/peterqiu0516/sub-store/internal/geoip"
)

// Flag detection rules: [regex, emoji, countryName, cityExtractor].
// Per review-resolution #16: Go RE2 supports \p{Regional_Indicator} and \uFE0F.
type geoRule struct {
	re            *regexp.Regexp
	emoji         string
	countryName   string
	countryNameCN string
	cityPatterns  []cityPattern
}

type cityPattern struct {
	re     *regexp.Regexp
	name   string
	nameCN string
}

var geoRules []geoRule

func init() {
	rules := []struct {
		pattern       string
		emoji         string
		countryName   string
		countryNameCN string
		cities        []struct{ pattern, name, nameCN string }
	}{
		{
			pattern:       `香港|港|hong\s*kong|\bhk\b`,
			emoji:         "🇭🇰",
			countryName:   "Hong Kong",
			countryNameCN: "香港",
			cities: []struct{ pattern, name, nameCN string }{
				{`香港|Hong\s*Kong|HK`, "Hong Kong", "香港"},
			},
		},
		{
			pattern:       `台湾|台灣|taiwan|\btw\b`,
			emoji:         "🇹🇼",
			countryName:   "Taiwan",
			countryNameCN: "台湾",
			cities: []struct{ pattern, name, nameCN string }{
				{`台北|臺北|Taipei`, "Taipei", "台北"},
				{`高雄|Gaoxiong|Kaohsiung`, "Kaohsiung", "高雄"},
			},
		},
		{
			pattern:       `新加坡|狮城|獅城|singapore|\bsg\b`,
			emoji:         "🇸🇬",
			countryName:   "Singapore",
			countryNameCN: "新加坡",
			cities: []struct{ pattern, name, nameCN string }{
				{`新加坡|Singapore|SG`, "Singapore", "新加坡"},
			},
		},
		{
			pattern:       `日本|东京|東京|大阪|japan|tokyo|osaka|\bjp\b`,
			emoji:         "🇯🇵",
			countryName:   "Japan",
			countryNameCN: "日本",
			cities: []struct{ pattern, name, nameCN string }{
				{`东京|東京|Tokyo`, "Tokyo", "东京"},
				{`大阪|Osaka`, "Osaka", "大阪"},
			},
		},
		{
			pattern:       `美国|美國|洛杉矶|洛杉磯|纽约|紐約|united\s*states|los\s*angeles|new\s*york|\bus\b|\busa\b`,
			emoji:         "🇺🇸",
			countryName:   "United States",
			countryNameCN: "美国",
			cities: []struct{ pattern, name, nameCN string }{
				{`洛杉矶|洛杉磯|Los\s*Angeles|LA`, "Los Angeles", "洛杉矶"},
				{`纽约|紐約|New\s*York|NYC`, "New York", "纽约"},
				{`硅谷|Silicon\s*Valley|San\s*Jose|SJC`, "San Jose", "硅谷"},
				{`西雅图|Seattle`, "Seattle", "西雅图"},
				{`芝加哥|Chicago`, "Chicago", "芝加哥"},
				{`达拉斯|Dallas`, "Dallas", "达拉斯"},
			},
		},
		{
			pattern:       `英国|英國|伦敦|倫敦|united\s*kingdom|london|\buk\b`,
			emoji:         "🇬🇧",
			countryName:   "United Kingdom",
			countryNameCN: "英国",
			cities: []struct{ pattern, name, nameCN string }{
				{`伦敦|倫敦|London`, "London", "伦敦"},
			},
		},
		{
			pattern:       `德国|德國|法兰克福|法蘭克福|germany|frankfurt|\bde\b`,
			emoji:         "🇩🇪",
			countryName:   "Germany",
			countryNameCN: "德国",
			cities: []struct{ pattern, name, nameCN string }{
				{`法兰克福|法蘭克福|Frankfurt`, "Frankfurt", "法兰克福"},
			},
		},
		{
			pattern:       `韩国|韓國|首尔|首爾|korea|seoul|\bkr\b`,
			emoji:         "🇰🇷",
			countryName:   "South Korea",
			countryNameCN: "韩国",
			cities: []struct{ pattern, name, nameCN string }{
				{`首尔|首爾|Seoul`, "Seoul", "首尔"},
			},
		},
	}
	for _, r := range rules {
		re, err := regexp.Compile("(?i)" + r.pattern)
		if err != nil {
			continue
		}
		var cps []cityPattern
		for _, c := range r.cities {
			cre, err := regexp.Compile("(?i)" + c.pattern)
			if err != nil {
				continue
			}
			cps = append(cps, cityPattern{re: cre, name: c.name, nameCN: c.nameCN})
		}
		geoRules = append(geoRules, geoRule{
			re:            re,
			emoji:         r.emoji,
			countryName:   r.countryName,
			countryNameCN: r.countryNameCN,
			cityPatterns:  cps,
		})
	}
}

// GeoInfo holds detected geographic information from a proxy name.
type GeoInfo struct {
	Flag        string // emoji flag
	CountryName string // English country name
	CountryCN   string // Chinese country name
	City        string // English city name (may be empty)
	CityCN      string // Chinese city name (may be empty)
}

// DetectFlag returns the flag emoji for a name, or 🏳️ if no match.
func DetectFlag(name string) string {
	text := strings.ToLower(name)
	for _, rule := range geoRules {
		if rule.re.MatchString(text) {
			return rule.emoji
		}
	}
	return "🏳️"
}

// DetectGeo returns full geographic info (flag + country + city) from a name.
// It tries GeoIP lookup on the server address first, falling back to
// name-based regex matching.
func DetectGeo(name string) GeoInfo {
	return detectGeoFromName(name)
}

// DetectGeoWithServer returns geographic info by looking up the proxy's
// server address via GeoIP first, falling back to name-based detection.
func DetectGeoWithServer(name, server string) GeoInfo {
	// Try GeoIP lookup first
	if server != "" {
		host := server
		// Strip port if present
		if h, _, err := net.SplitHostPort(server); err == nil {
			host = h
		}
		if result := geoip.LookupHost(host); result != nil {
			return GeoInfo{
				Flag:        geoip.CountryEmoji(result.CountryCode),
				CountryName: result.CountryName,
				CountryCN:   result.CountryNameCN,
				City:        result.CityName,
				CityCN:      result.CityNameCN,
			}
		}
	}
	// Fallback to name-based detection
	return detectGeoFromName(name)
}

// detectGeoFromName does name-based regex matching for geo detection.
func detectGeoFromName(name string) GeoInfo {
	text := strings.ToLower(name)
	for _, rule := range geoRules {
		if rule.re.MatchString(text) {
			city := ""
			cityCN := ""
			for _, cp := range rule.cityPatterns {
				if cp.re.MatchString(text) {
					city = cp.name
					cityCN = cp.nameCN
					break
				}
			}
			return GeoInfo{
				Flag:        rule.emoji,
				CountryName: rule.countryName,
				CountryCN:   rule.countryNameCN,
				City:        city,
				CityCN:      cityCN,
			}
		}
	}
	return GeoInfo{}
}

// removeFlagRe matches leading flag emoji sequences and whitespace.
// Regional Indicator letters are U+1F1E6–U+1F1FF; ZWJ is U+200D; VS16 is U+FE0F.
var removeFlagRe = regexp.MustCompile(`^[\x{1F1E6}-\x{1F1FF}\x{FE0F}\x{200D}\s]+`)

// RemoveFlag strips leading flag emoji and whitespace from a name.
// Per review-resolution #16: handles ZWJ (\u200D) and variation selector (\uFE0F).
func RemoveFlag(name string) string {
	cleaned := removeFlagRe.ReplaceAllString(name, "")
	// Also strip leading 🏳️
	cleaned = strings.TrimLeft(cleaned, "🏳️ ")
	return strings.TrimSpace(cleaned)
}

// NormalizeTaiwanFlag maps the Taiwan flag to different flags based on `tw` mode.
// Per review-resolution #19: three modes — ws/tw/default.
func NormalizeTaiwanFlag(flag, mode string) string {
	if flag != "🇹🇼" {
		return flag
	}
	switch mode {
	case "ws":
		return "🇼🇸"
	case "tw":
		return "🇹🇼"
	default:
		return "🇨🇳"
	}
}

// IsASCII checks if a string contains only ASCII characters.
func IsASCII(s string) bool {
	for i := 0; i < len(s); i++ {
		if s[i] > 127 {
			return false
		}
	}
	return true
}