pancakeswapnft/squad.go

package pancakeswapnft

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"sort"
	"strconv"
	"sync"
	"time"

	"github.com/pkg/errors"
	"github.com/rs/zerolog/log"
)

type Squad struct {
	TokenId     int64  `json:"tokenId,string"`
	Name        string `json:"name"`
	Description string `json:"description"`
	Image       struct {
		Original  string `json:"original"`
		Thumbnail string `json:"thumbnail"`
		Mp4       string `json:"mp4"`
		Webm      string `json:"webm"`
		Gif       string `json:"gif"`
	} `json:"image"`
	CreatedAt  time.Time `json:"createdAt"`
	UpdatedAt  time.Time `json:"updatedAt"`
	Attributes []struct {
		TraitType   string `json:"traitType"`
		Value       string `json:"value"`
		DisplayType string `json:"displayType"`
	} `json:"attributes"`
	Collection struct {
		Name string `json:"name"`
	} `json:"collection"`
	RarityScore float64 `json:"rarity_score"`
	Rank        int     `json:"rank"`
	TraitCount  int     `json:"trait_count"`
}

type SquadService struct {
	Collection Collection
	squads     map[int64]Squad
	squadslock *sync.RWMutex
}

func NewSquadService(collection Collection) *SquadService {
	return &SquadService{
		Collection: collection,
		squads:     make(map[int64]Squad),
		squadslock: &sync.RWMutex{},
	}
}

func (s *SquadService) LoadCacheFromFile(filename string) error {
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		err := s.LoadCacheFromWeb(s.Collection.FirstTokenId, s.Collection.TokenCount)
		if err != nil {
			return errors.Wrap(err, "loading cache from web")
		}
		err = s.SaveCacheToFile(filename)
		if err != nil {
			return errors.Wrap(err, "saving cache to file")
		}
		return nil
	}
	f, err := os.OpenFile(filename, os.O_RDONLY, os.ModePerm)
	if err != nil {
		return errors.Wrapf(err, "opening file %s", filename)
	}
	var res map[int64]Squad
	err = json.NewDecoder(f).Decode(&res)
	if err != nil {
		return errors.Wrap(err, "decoding json")
	}
	s.squadslock.Lock()
	s.squads = res
	defer s.squadslock.Unlock()
	return nil
}

func (s *SquadService) SaveCacheToFile(filepath string) error {
	file, _ := os.OpenFile(filepath, os.O_CREATE, os.ModePerm)
	defer file.Close()
	s.squadslock.RLock()
	defer s.squadslock.RUnlock()
	err := json.NewEncoder(file).Encode(s.squads)
	if err != nil {
		return errors.Wrapf(err, "encoding to json & writing to file %s", filepath)
	}
	return nil
}

func (s *SquadService) getSquadFromWeb(id int) (Squad, error) {
	resp, err := http.Get(fmt.Sprintf("https://nft.pancakeswap.com/api/v1/collections/%s/tokens/%d", s.Collection.Id, id))
	if err != nil {
		return Squad{}, err
	}
	defer resp.Body.Close()

	var res struct {
		Data Squad `json:"data"`
	}
	err = json.NewDecoder(resp.Body).Decode(&res)
	if err != nil {
		return Squad{}, err
	}
	return res.Data, nil
}

func (s *SquadService) getRarityData() (map[string]map[string]int, error) {
	resp, err := http.Get(fmt.Sprintf("https://nft.pancakeswap.com/api/v1/collections/%s/distribution", s.Collection.Id))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var res struct {
		Data map[string]map[string]int `json:"data"`
	}
	err = json.NewDecoder(resp.Body).Decode(&res)
	if err != nil {
		return nil, err
	}
	return res.Data, nil
}

func (s *SquadService) LoadCacheFromWeb(starting int, count int) error {
	squadcount := 10000
	var squads []*Squad
	for i := starting; i <= count+starting-1; i++ {
		log.Printf("getting %d", i)
		s, err := s.getSquadFromWeb(i)
		if err != nil {
			return errors.Wrapf(err, "getting squad %d", i)
		}
		squads = append(squads, &s)
	}

	raritydata, err := s.getRarityData()
	if err != nil {
		return errors.Wrap(err, "getting rarity data from api")
	}

	raritydata["traitcount"] = make(map[string]int)
	for _, squad := range squads {
		traitcount := 0
		for _, attribute := range squad.Attributes {
			if attribute.Value != "None" {
				traitcount++
			}
		}
		raritydata["traitcount"][strconv.Itoa(traitcount)]++
		squad.TraitCount = traitcount
		if traitcount == 0 {
			squadcount--
		}
	}

	for _, squad := range squads {
		rarityScore := float64(0)
		for _, attribute := range squad.Attributes {
			rarityScore += 1 / (float64(raritydata[attribute.TraitType][attribute.Value]) / float64(squadcount))
		}
		rarityScore += 1 / (float64(raritydata["traitcount"][strconv.Itoa(squad.TraitCount)]) / float64(squadcount))
		squad.RarityScore = rarityScore
	}

	sort.Slice(squads, func(i, j int) bool {
		return squads[i].RarityScore > squads[j].RarityScore
	})

	for i, squad := range squads {
		squad.Rank = i + 1
	}
	s.squadslock.Lock()
	defer s.squadslock.Unlock()
	for _, squad := range squads {
		s.squads[squad.TokenId] = *squad
	}
	return nil
}

func (s *SquadService) GetSquad(id int64) (Squad, error) {
	s.squadslock.RLock()
	defer s.squadslock.RUnlock()
	if s, ok := s.squads[id]; !ok {
		return Squad{}, errors.Errorf("squad %d not found", id)
	} else {
		return s, nil
	}
}

func (s *Squad) GetScoreByPrice(floorService *FloorService, price float64) float64 {
	return 1 + (4 / ((price - (floorService.GetFloor() - 1.5)) * float64(s.Rank) / float64(10000)))
}

func (s *Squad) GetPriceByScore(floorService *FloorService, score float64) float64 {
	return (floorService.GetFloor() - 1.5) + (4 / ((score - 1) * float64(s.Rank) / float64(10000)))
}