github.com/nxtrace/NTrace-core@v1.3.1-0.20240513132635-39169291e8c9/dn42/ptr.go (about)

     1  package dn42
     2  
     3  import (
     4  	"encoding/csv"
     5  	"errors"
     6  	"fmt"
     7  	"os"
     8  	"regexp"
     9  	"strings"
    10  
    11  	"github.com/spf13/viper"
    12  )
    13  
    14  type PtrRow struct {
    15  	IATACode string
    16  	LtdCode  string
    17  	Region   string
    18  	City     string
    19  }
    20  
    21  func matchesPattern(prefix string, s string) bool {
    22  	pattern := fmt.Sprintf(`^(.*[-.\d]|^)%s[-.\d].*$`, prefix)
    23  
    24  	r, err := regexp.Compile(pattern)
    25  	if err != nil {
    26  		fmt.Println("Invalid regular expression:", err)
    27  		return false
    28  	}
    29  
    30  	return r.MatchString(s)
    31  }
    32  
    33  func FindPtrRecord(ptr string) (PtrRow, error) {
    34  	path := viper.Get("ptrPath").(string)
    35  	f, err := os.Open(path)
    36  	if err != nil {
    37  		return PtrRow{}, err
    38  	}
    39  	defer f.Close()
    40  
    41  	r := csv.NewReader(f)
    42  	rows, err := r.ReadAll()
    43  	if err != nil {
    44  		return PtrRow{}, err
    45  	}
    46  	// 转小写
    47  	ptr = strings.ToLower(ptr)
    48  	// 先查城市名
    49  	for _, row := range rows {
    50  		city := row[3]
    51  		if city == "" {
    52  			continue
    53  		}
    54  		city = strings.ReplaceAll(city, " ", "")
    55  		city = strings.ToLower(city)
    56  
    57  		if matchesPattern(city, ptr) {
    58  			return PtrRow{
    59  				LtdCode: row[1],
    60  				Region:  row[2],
    61  				City:    row[3],
    62  			}, nil
    63  		}
    64  	}
    65  	// 查 IATA Code
    66  	for _, row := range rows {
    67  		iata := row[0]
    68  		if iata == "" {
    69  			continue
    70  		}
    71  		iata = strings.ToLower(iata)
    72  		if matchesPattern(iata, ptr) {
    73  			return PtrRow{
    74  				IATACode: iata,
    75  				LtdCode:  row[1],
    76  				Region:   row[2],
    77  				City:     row[3],
    78  			}, nil
    79  		}
    80  	}
    81  
    82  	return PtrRow{}, errors.New("ptr not found")
    83  }