// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.

package nodes

import (
	"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/iplibrary"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/TeaOSLab/EdgeDNS/internal/agents"
	"github.com/TeaOSLab/EdgeDNS/internal/dbs"
	"github.com/TeaOSLab/EdgeDNS/internal/models"
	"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
	"github.com/TeaOSLab/EdgeDNS/internal/rpc"
	"github.com/iwind/TeaGo/lists"
	"github.com/iwind/TeaGo/types"
	"net"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// RouteManager 线路管理器
type RouteManager struct {
	allRouteMap  map[int64]*models.NSRoute // routeId => route
	userRouteMap map[int64][]int64         // userId => sorted routeIds

	db      *dbs.DB
	version int64
	locker  sync.RWMutex

	notifier chan bool
	readyCh  chan struct{} // 初始加载完成后关闭

	ispRouteMap   map[string]string // name => code
	chinaRouteMap map[string]string // name => code
	worldRouteMap map[string]string // name => code
}

// NewRouteManager 获取新线路管理器对象
func NewRouteManager(db *dbs.DB) *RouteManager {
	return &RouteManager{
		db:           db,
		allRouteMap:  map[int64]*models.NSRoute{},
		userRouteMap: map[int64][]int64{},

		notifier: make(chan bool, 8),
		readyCh:  make(chan struct{}),

		ispRouteMap:   map[string]string{},
		chinaRouteMap: map[string]string{},
		worldRouteMap: map[string]string{},
	}
}

// Start 启动自动任务
func (this *RouteManager) Start() {
	remotelogs.Println("ROUTE_MANAGER", "starting ...")

	// 初始化公共线路
	this.loadDefaultRoutes()

	// 从本地数据库中加载数据
	err := this.Load()
	if err != nil {
		if rpc.IsConnError(err) {
			remotelogs.Debug("ROUTE_MANAGER", "load failed: "+err.Error())
		} else {
			remotelogs.Error("ROUTE_MANAGER", "load failed: "+err.Error())
		}
	}

	// 初始化运行
	err = this.LoopAll()
	if err != nil {
		if rpc.IsConnError(err) {
			remotelogs.Debug("ROUTE_MANAGER", "loop failed: "+err.Error())
		} else {
			remotelogs.Error("ROUTE_MANAGER", "loop failed: "+err.Error())
		}
	}

	// 通知初始加载完成
	close(this.readyCh)

	// 更新
	var ticker = time.NewTicker(1 * time.Minute)
	for {
		select {
		case <-ticker.C:
		case <-this.notifier:
		}

		err := this.LoopAll()
		if err != nil {
			if rpc.IsConnError(err) {
				remotelogs.Debug("ROUTE_MANAGER", "loop failed: "+err.Error())
			} else {
				remotelogs.Error("ROUTE_MANAGER", "loop failed: "+err.Error())
			}
		}
	}
}

// Load 从数据库中加载数据
func (this *RouteManager) Load() error {
	var offset int64 = 0
	var size int64 = 10000
	for {
		routes, err := this.db.ListRoutes(offset, size)
		if err != nil {
			return err
		}

		if len(routes) == 0 {
			break
		}

		this.locker.Lock()
		for _, route := range routes {
			this.addRoute(route)

			if route.Version > this.version {
				this.version = route.Version
			}
		}
		this.locker.Unlock()

		offset += size
	}

	if this.version > 0 {
		this.version++
	}

	return nil
}

func (this *RouteManager) LoopAll() error {
	for {
		hasNext, err := this.Loop()
		if err != nil {
			return err
		}
		if !hasNext {
			break
		}
	}
	return nil
}

// Loop 单次循环任务
func (this *RouteManager) Loop() (hasNext bool, err error) {
	client, err := rpc.SharedRPC()
	if err != nil {
		return false, err
	}
	resp, err := client.NSRouteRPC.ListNSRoutesAfterVersion(client.Context(), &pb.ListNSRoutesAfterVersionRequest{
		Version: this.version,
		Size:    20000,
	})
	if err != nil {
		return false, err
	}
	var routes = resp.NsRoutes
	if len(routes) == 0 {
		return false, nil
	}
	for _, route := range routes {
		this.processRoute(route)

		if route.Version > this.version {
			this.version = route.Version
		}
	}
	this.version++
	return true, nil
}

// FindRouteCodes 查找一个地址对应的线路
func (this *RouteManager) FindRouteCodes(ip string, domainUserId int64) (result []string) {
	var netIP = net.ParseIP(ip)
	if len(netIP) == 0 {
		return nil
	}

	// 自定义route
	this.locker.RLock()

	// 先查找用户自定义的
	if domainUserId > 0 {
		var userRouteIds = this.userRouteMap[domainUserId]
		for _, routeId := range userRouteIds {
			route, ok := this.allRouteMap[routeId]
			if ok && route.Contains(netIP) {
				result = append(result, route.RealCode())
			}
		}
	}

	// 再查找公共的
	var publicRouteIds = this.userRouteMap[0]
	for _, routeId := range publicRouteIds {
		route, ok := this.allRouteMap[routeId]
		if ok && route.Contains(netIP) {
			result = append(result, route.RealCode())
		}
	}

	this.locker.RUnlock()

	// 解析公用线路
	var ipResult = iplibrary.LookupIP(ip)
	if ipResult != nil && ipResult.IsOk() {
		// 运营商
		for _, providerCode := range ipResult.ProviderCodes() {
			code, ok := this.ispRouteMap[providerCode]
			if ok {
				result = append(result, code)

				// 单次只能有一个匹配
				break
			}
		}

		// 省|州|城市
		if ipResult.ProvinceId() > 0 {
			result = append(result, "region:province:"+types.String(ipResult.ProvinceId()))
		}
		if ipResult.CityId() > 0 {
			result = append(result, "region:city:"+types.String(ipResult.CityId()))
		}
		if ipResult.TownId() > 0 {
			result = append(result, "region:town:"+types.String(ipResult.TownId()))
		}

		// 中国省市
		for _, provinceCode := range ipResult.ProvinceCodes() {
			// 中国
			code, ok := this.chinaRouteMap[provinceCode]
			if ok {
				result = append(result, code)

				// 兼容以前的拼写错误
				switch code {
				case "china:province:hebei":
					result = append(result, "china:province:heibei")
				case "china:province:heibei":
					result = append(result, "china:province:hebei")
				case "china:jilin":
					result = append(result, "china:province:jilin")
				case "china:province:jilin":
					result = append(result, "china:jilin")
				}

				// 香港
				switch code {
				case dnsconfigs.ChinaProvinceCodeHK:
					result = append(result, dnsconfigs.WorldRegionCodeHK, dnsconfigs.WorldRegionCodeChinaAbroad)

				// 澳门
				case dnsconfigs.ChinaProvinceCodeMO:
					result = append(result, dnsconfigs.WorldRegionCodeMO, dnsconfigs.WorldRegionCodeChinaAbroad)

				// 台湾
				case dnsconfigs.ChinaProvinceCodeTW:
					result = append(result, dnsconfigs.WorldRegionCodeTW, dnsconfigs.WorldRegionCodeChinaAbroad)
				default:
					result = append(result, dnsconfigs.WorldRegionCodeChinaMainland)
				}

				// 单次只能有一个匹配
				break
			}
		}

		// 国家/地区
		for _, countryCode := range ipResult.CountryCodes() {
			code, ok := this.worldRouteMap[countryCode]
			if ok {
				// 中国全境（world:CN）必须优先于「海外」匹配，否则前面中国省市若误判为 HK/MO/TW 已加入 world:CN:abroad 时会先命中海外线
				if code == dnsconfigs.WorldRegionCodeChina {
					result = append([]string{code}, result...)
				} else {
					result = append(result, code)
					// 中国海外
					if code != dnsconfigs.WorldRegionCodeChina {
						result = append(result, dnsconfigs.WorldRegionCodeChinaAbroad)
					}
				}

				// 单次只能有一个匹配
				break
			}
		}
	}

	// 搜索引擎线路
	if agents.SharedManager != nil {
		var agentCode = agents.SharedManager.LookupIP(ip)
		if len(agentCode) > 0 {
			result = append(result, "agent:"+agentCode, "agent" /** 所有搜索引擎 **/)
		}
	}

	return
}

// NotifyUpdate 通知更新
func (this *RouteManager) NotifyUpdate() {
	select {
	case this.notifier <- true:
	default:
	}
}

func (this *RouteManager) loadDefaultRoutes() {
	for _, route := range dnsconfigs.AllDefaultISPRoutes {
		for _, name := range route.AliasNames {
			this.ispRouteMap[name] = route.Code
		}
	}
	for _, route := range dnsconfigs.AllDefaultChinaProvinceRoutes {
		for _, name := range route.AliasNames {
			this.chinaRouteMap[name] = route.Code
		}
	}
	for _, route := range dnsconfigs.AllDefaultWorldRegionRoutes {
		for _, name := range route.AliasNames {
			this.worldRouteMap[name] = route.Code
		}
		// 用线路 Code 中的国家/地区 ISO 码做映射，使 IP 库返回的 ISO 码（如 US、CN、HK）能命中对应线路
		if strings.HasPrefix(route.Code, "world:") {
			parts := strings.Split(route.Code, ":")
			if len(parts) == 2 && len(parts[1]) == 2 {
				this.worldRouteMap[parts[1]] = route.Code
			}
			if len(parts) == 3 && len(parts[2]) == 2 {
				this.worldRouteMap[strings.ToUpper(parts[2])] = route.Code
			}
		}
	}
}

// 添加线路
func (this *RouteManager) addRoute(route *models.NSRoute) {
	// 不需要加锁，因为此函数均在锁内调用

	// 从老的用户中删除
	oldRoute, ok := this.allRouteMap[route.Id]
	if ok {
		var oldUserId = oldRoute.UserId
		if oldUserId != route.UserId {
			userRouteIds, ok := this.userRouteMap[oldUserId]
			if ok {
				this.userRouteMap[oldUserId] = this.removeId(userRouteIds, route.Id)
				if len(userRouteIds) == 0 {
					delete(this.userRouteMap, oldUserId)
				}
			}
		}
	}

	// 添加
	this.allRouteMap[route.Id] = route

	userRouteIds, ok := this.userRouteMap[route.UserId]
	if ok {
		// 重新按优先级、排序、ID排序
		var userRoutes = []*models.NSRoute{}
		for _, userRouteId := range userRouteIds {
			userRoute, ok := this.allRouteMap[userRouteId]
			if ok {
				userRoutes = append(userRoutes, userRoute)
			}
		}
		if !lists.ContainsInt64(userRouteIds, route.Id) {
			userRoutes = append(userRoutes, route)
		}

		sort.Slice(userRoutes, func(i, j int) bool {
			var userRoute1 = userRoutes[i]
			var userRoute2 = userRoutes[j]
			if userRoute1.Priority != userRoute2.Priority {
				return userRoute1.Priority > userRoute2.Priority
			}
			if userRoute1.Order != userRoute2.Order {
				return userRoute1.Order > userRoute2.Order
			}
			return userRoute1.Id < userRoute2.Id
		})

		var newUserRouteIds = []int64{}
		for _, userRoute := range userRoutes {
			newUserRouteIds = append(newUserRouteIds, userRoute.Id)
		}
		this.userRouteMap[route.UserId] = newUserRouteIds
	} else {
		this.userRouteMap[route.UserId] = []int64{route.Id}
	}
}

// 删除线路
func (this *RouteManager) removePBRoute(route *pb.NSRoute) {
	delete(this.allRouteMap, route.Id)
	userRouteIds, ok := this.userRouteMap[route.UserId]
	if ok {
		userRouteIds = this.removeId(userRouteIds, route.Id)
		if len(userRouteIds) == 0 {
			delete(this.userRouteMap, route.UserId)
		} else {
			this.userRouteMap[route.UserId] = userRouteIds
		}
	}
}

// 处理线路
func (this *RouteManager) processRoute(route *pb.NSRoute) {
	if !route.IsOn || route.IsDeleted {
		this.locker.Lock()
		this.removePBRoute(route)
		this.locker.Unlock()

		// 从数据库中删除
		if this.db != nil {
			err := this.db.DeleteRoute(route.Id)
			if err != nil {
				remotelogs.Error("ROUTE_MANAGER", "delete route from db failed: "+err.Error())
			}
		}

		return
	}

	// 存入数据库
	if this.db != nil {
		exists, err := this.db.ExistsRoute(route.Id)
		if err != nil {
			remotelogs.Error("ROUTE_MANAGER", "query failed: "+err.Error())
		} else {
			if exists {
				err = this.db.UpdateRoute(route.Id, route.UserId, route.RangesJSON, route.Order, route.Priority, route.Version)
				if err != nil {
					remotelogs.Error("ROUTE_MANAGER", "update failed: "+err.Error())
				}
			} else {
				err = this.db.InsertRoute(route.Id, route.UserId, route.RangesJSON, route.Order, route.Priority, route.Version)
				if err != nil {
					remotelogs.Error("ROUTE_MANAGER", "insert failed: "+err.Error())
				}
			}
		}
	}

	ranges, err := models.InitRangesFromJSON(route.RangesJSON)
	if err != nil {
		remotelogs.Error("ROUTE_MANAGER", "decode routes '"+strconv.FormatInt(route.Id, 10)+"' failed: "+err.Error())
		return
	}

	var nsRoute = &models.NSRoute{
		Id:       route.Id,
		Ranges:   ranges,
		Priority: route.Priority,
		Order:    route.Order,
		UserId:   route.UserId,
		Version:  route.Version,
	}

	this.locker.Lock()
	this.addRoute(nsRoute)
	this.locker.Unlock()
}

// 从一组ID中删除某个ID
func (this *RouteManager) removeId(ids []int64, id int64) []int64 {
	var result = []int64{}
	for _, id2 := range ids {
		if id2 == id {
			continue
		}
		result = append(result, id2)
	}
	return result
}