package store

import (
	"encoding/json"
	"errors"
	"time"

	badger "github.com/dgraph-io/badger/v4"
)

type Entry struct {
	TokenIDs  []string  `json:"token_ids"`
	UpdatedAt time.Time `json:"updated_at"`
}

// ContractEntry stores all token ownership data for a contract
type ContractEntry struct {
	// Map of tokenId -> owner address (canonical)
	TokenOwners map[string]string `json:"token_owners"`
	UpdatedAt   time.Time         `json:"updated_at"`
	MaxScanned  int              `json:"max_scanned"`
}

type Cache struct {
	db *badger.DB
}

func New(dir string) (*Cache, error) {
	opts := badger.DefaultOptions(dir)
	opts.Logger = nil
	db, err := badger.Open(opts)
	if err != nil { return nil, err }
	return &Cache{db: db}, nil
}

func (c *Cache) Close() error { return c.db.Close() }

func (c *Cache) Get(key string) (*Entry, bool) {
	var e Entry
	err := c.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(key))
		if err != nil { return err }
		return item.Value(func(v []byte) error { return json.Unmarshal(v, &e) })
	})
	if err != nil { return nil, false }
	return &e, true
}

func (c *Cache) Set(key string, ids []string) error {
	val, _ := json.Marshal(Entry{TokenIDs: ids, UpdatedAt: time.Now().UTC()})
	return c.db.Update(func(txn *badger.Txn) error {
		return txn.Set([]byte(key), val)
	})
}

// SetContract stores all token ownership data for a contract
func (c *Cache) SetContract(key string, tokenOwners map[string]string, maxScanned int) error {
	val, _ := json.Marshal(ContractEntry{
		TokenOwners: tokenOwners,
		UpdatedAt:   time.Now().UTC(),
		MaxScanned:  maxScanned,
	})
	return c.db.Update(func(txn *badger.Txn) error {
		return txn.Set([]byte(key), val)
	})
}

// GetContract retrieves contract ownership data
func (c *Cache) GetContract(key string) (*ContractEntry, bool) {
	var e ContractEntry
	err := c.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(key))
		if err != nil { return err }
		return item.Value(func(v []byte) error { return json.Unmarshal(v, &e) })
	})
	if err != nil { return nil, false }
	return &e, true
}

func (c *Cache) Delete(key string) error {
	return c.db.Update(func(txn *badger.Txn) error {
		err := txn.Delete([]byte(key))
		if errors.Is(err, badger.ErrKeyNotFound) { return nil }
		return err
	})
}

// DeleteContract removes contract cache entry
func (c *Cache) DeleteContract(key string) error {
	return c.Delete(key)
}