package tokenencoding

import (
	"embed"
	"encoding/base64"
	"fmt"
	"path"
	"strconv"
	"strings"
	"sync"

	tiktoken "github.com/pkoukk/tiktoken-go"
)

//go:embed assets/*.tiktoken
var bpeAssets embed.FS

var configureOnce sync.Once

func GetEncoding(name string) (*tiktoken.Tiktoken, error) {
	configureOnce.Do(func() {
		tiktoken.SetBpeLoader(offlineLoader{})
	})
	return tiktoken.GetEncoding(name)
}

type offlineLoader struct{}

func (offlineLoader) LoadTiktokenBpe(tiktokenBpeFile string) (map[string]int, error) {
	contents, err := bpeAssets.ReadFile("assets/" + path.Base(tiktokenBpeFile))
	if err != nil {
		return nil, err
	}

	bpeRanks := make(map[string]int)
	for _, line := range strings.Split(string(contents), "\n") {
		if line == "" {
			continue
		}
		parts := strings.Split(line, " ")
		if len(parts) != 2 {
			return nil, fmt.Errorf("invalid tiktoken BPE line %q", line)
		}
		token, err := base64.StdEncoding.DecodeString(parts[0])
		if err != nil {
			return nil, err
		}
		rank, err := strconv.Atoi(parts[1])
		if err != nil {
			return nil, err
		}
		bpeRanks[string(token)] = rank
	}
	return bpeRanks, nil
}