package shreder

import (
	"context"
	"fmt"
	"sync"

	"github.com/gagliardetto/solana-go"
	"github.com/gagliardetto/solana-go/rpc"
	lru "github.com/hashicorp/golang-lru/v2"
	"github.com/panjf2000/ants/v2"
)

type TableInfo struct {
	overErrCount int

	addresses []solana.PublicKey
}

type AddressTables struct {
	showTableLoaded bool

	rpcClient *rpc.Client
	loadMux   sync.Mutex
	tables    *lru.Cache[solana.PublicKey, *TableInfo]
	loading   map[solana.PublicKey]struct{}

	pool *ants.Pool
}

func NewAddressTables(rpcClient *rpc.Client, showTableLoaded bool) *AddressTables {
	pool, _ := ants.NewPool(5, ants.WithPreAlloc(true), ants.WithNonblocking(true))
	cache, _ := lru.New[solana.PublicKey, *TableInfo](10000)
	return &AddressTables{
		rpcClient: rpcClient,
		tables:    cache,
		loading:   make(map[solana.PublicKey]struct{}),
		pool:      pool,

		showTableLoaded: showTableLoaded,
	}
}

func (at *AddressTables) loadAddressTable(tablePubkey solana.PublicKey) ([]solana.PublicKey, error) {
	// decode acc
	acc, err := at.rpcClient.GetAccountInfoWithOpts(context.Background(), tablePubkey, &rpc.GetAccountInfoOpts{
		Encoding: solana.EncodingBase64,
	})
	if err != nil {
		return nil, err
	}
	data := acc.GetBinary()
	if len(data) <= 56 {
		return nil, fmt.Errorf("account data too short")
	}
	offset := 56
	var addresses solana.PublicKeySlice = make([]solana.PublicKey, 0, (len(data)-offset)/32)
	for offset+32 <= len(data) {
		addresses = append(addresses, solana.PublicKeyFromBytes(data[offset:offset+32]))
		offset += 32
	}
	// addresses = append(addresses, solana.PublicKeyFromBytes(data[start:start+32]))
	return addresses, nil

}
func (at *AddressTables) load(tablePubkey solana.PublicKey) {
	_ = at.pool.Submit(func() {
		at.loadMux.Lock()
		_, loading := at.loading[tablePubkey]
		if loading {
			at.loadMux.Unlock()
			return
		}
		at.loading[tablePubkey] = struct{}{}
		at.loadMux.Unlock()

		table, err := at.loadAddressTable(tablePubkey)
		if err != nil {
			logger.Error("loadAddressTable failed", "err", err, "table", tablePubkey)
			at.loadMux.Lock()
			delete(at.loading, tablePubkey)
			at.loadMux.Unlock()
			return
		}
		at.loadMux.Lock()
		delete(at.loading, tablePubkey)
		at.loadMux.Unlock()

		at.tables.Add(tablePubkey, &TableInfo{
			addresses: table,
		})
		if at.showTableLoaded {
			total := at.tables.Len()
			logger.Info("loadAddressTable", "table", tablePubkey.String(), "table count:", total)
		}
	})
}

func (at *AddressTables) FillToTx(tx *versionedTransaction, tablePubkey solana.PublicKey, idx []uint8) bool {
	addresses, ok := at.tables.Get(tablePubkey)
	if !ok {
		at.load(tablePubkey)
		return false
	}

	for _, i := range idx {
		if int(i) >= len(addresses.addresses) {
			logger.Error("over loadAddressTable failed", "idx", i, "table", tablePubkey)
			addresses.overErrCount++
			if addresses.overErrCount > 10 {
				at.load(tablePubkey)
			}
			return false
		}
		tx.Message.StaticAccountKeys = append(tx.Message.StaticAccountKeys, addresses.addresses[i])
	}
	return true
}

func (at *AddressTables) GetAddressTable(tablePubkey solana.PublicKey, idx []uint8) []solana.PublicKey {
	addresses, ok := at.tables.Get(tablePubkey)
	if !ok {
		at.load(tablePubkey)
		return nil
	}

	var result solana.PublicKeySlice = make([]solana.PublicKey, 0, len(idx))
	for _, i := range idx {
		if int(i) >= len(addresses.addresses) {
			logger.Error("over loadAddressTable failed", "idx", i, "table", tablePubkey)
			addresses.overErrCount++
			if addresses.overErrCount > 10 {
				at.load(tablePubkey)
			}
			break
		}
		result = append(result, addresses.addresses[i])
	}
	return result
}