package pd

import (
	"encoding/json"
	"fmt"
)

type PD struct{}

func (p *PD) ListSchedulers() ([]string, error) {
	resp, err := pdController.GetDriver().ShowSchedulers()
	if err != nil {
		return nil, err
	}
	var schedulers []string
	err = json.Unmarshal(resp, &schedulers)
	return schedulers, err
}

func (p *PD) ListStores() ([]StoreInfo, error) {
	resp, err := pdController.GetDriver().ShowStores()
	if err != nil {
		return nil, err
	}
	var storesInfo StoresInfo
	err = json.Unmarshal(resp, &storesInfo)
	if err != nil {
		return nil, err
	}
	return storesInfo.Stores, nil
}

func (p *PD) ListRegions() ([]Region, error) {
	resp, err := pdController.GetDriver().ShowRegions()
	if err != nil {
		return nil, err
	}
	var regionsInfo RegionsInfo
	err = json.Unmarshal(resp, &regionsInfo)
	if err != nil {
		return nil, err
	}
	return regionsInfo.Regions, nil
}

func (p *PD) RemoveScheduler(name string) error {
	return pdController.GetDriver().RemoveScheduler(name)
}

func (p *PD) ExistsStoreScheduler(name string, storeID uint64) (bool, error) {
	schedulers, err := p.ListSchedulers()
	if err != nil {
		return false, err
	}
	for _, s := range schedulers {
		if string(s) == fmt.Sprintf("%s-%d", name, storeID) {
			return true, nil
		}
	}
	return false, nil
}

func (p *PD) GetStoreRegionMap() (map[uint64][]Region, error) {
	return p.BuildStoreRegionMap(func(s StoreInfo) ([]Region, error) {
		return s.GetRegions()
	})
}

func (p *PD) GetStoreRegionLeaderMap() (map[uint64][]Region, error) {
	return p.BuildStoreRegionMap(func(s StoreInfo) ([]Region, error) {
		return s.GetLeaders()
	})
}

func (p *PD) BuildStoreRegionMap(resolve func(StoreInfo) ([]Region, error)) (map[uint64][]Region, error) {
	stores, err := p.ListStores()
	if err != nil {
		return nil, err
	}
	m := make(map[uint64][]Region)
	for _, s := range stores {
		if _, ok := m[s.Store.ID]; !ok {
			regions, err := resolve(s)
			if err != nil {
				return nil, err
			}
			m[s.Store.ID] = regions
		}
	}
	return m, nil
}