package storage

import (
	"context"

	"github.com/filecoin-project/go-address"
	"github.com/filecoin-project/go-state-types/abi"
	"github.com/filecoin-project/go-state-types/big"

	"github.com/filecoin-project/lotus/api"
	"github.com/filecoin-project/lotus/chain/actors/builtin/miner"
	"github.com/filecoin-project/lotus/chain/types"
)

type addrSelectApi interface {
	WalletBalance(context.Context, address.Address) (types.BigInt, error)
	WalletHas(context.Context, address.Address) (bool, error)

	StateAccountKey(context.Context, address.Address, types.TipSetKey) (address.Address, error)
	StateLookupID(context.Context, address.Address, types.TipSetKey) (address.Address, error)
}

type AddressSelector struct {
	api.AddressConfig
}

func (as *AddressSelector) AddressFor(ctx context.Context, a addrSelectApi, mi miner.MinerInfo, use api.AddrUse, goodFunds, minFunds abi.TokenAmount) (address.Address, abi.TokenAmount, error) {
	if as == nil {
		// should only happen in some tests
		log.Warnw("smart address selection disabled, using worker address")
		return mi.Worker, big.Zero(), nil
	}

	var addrs []address.Address
	switch use {
	case api.PreCommitAddr:
		addrs = append(addrs, as.PreCommitControl...)
	case api.CommitAddr:
		addrs = append(addrs, as.CommitControl...)
	case api.TerminateSectorsAddr:
		addrs = append(addrs, as.TerminateControl...)
	case api.DealPublishAddr:
		addrs = append(addrs, as.DealPublishControl...)
	default:
		defaultCtl := map[address.Address]struct{}{}
		for _, a := range mi.ControlAddresses {
			defaultCtl[a] = struct{}{}
		}
		delete(defaultCtl, mi.Owner)
		delete(defaultCtl, mi.Worker)

		configCtl := append([]address.Address{}, as.PreCommitControl...)
		configCtl = append(configCtl, as.CommitControl...)
		configCtl = append(configCtl, as.TerminateControl...)
		configCtl = append(configCtl, as.DealPublishControl...)

		for _, addr := range configCtl {
			if addr.Protocol() != address.ID {
				var err error
				addr, err = a.StateLookupID(ctx, addr, types.EmptyTSK)
				if err != nil {
					log.Warnw("looking up control address", "address", addr, "error", err)
					continue
				}
			}

			delete(defaultCtl, addr)
		}

		for a := range defaultCtl {
			addrs = append(addrs, a)
		}
	}

	if len(addrs) == 0 || !as.DisableWorkerFallback {
		addrs = append(addrs, mi.Worker)
	}
	if !as.DisableOwnerFallback {
		addrs = append(addrs, mi.Owner)
	}

	return pickAddress(ctx, a, mi, goodFunds, minFunds, addrs)
}

func pickAddress(ctx context.Context, a addrSelectApi, mi miner.MinerInfo, goodFunds, minFunds abi.TokenAmount, addrs []address.Address) (address.Address, abi.TokenAmount, error) {
	leastBad := mi.Worker
	bestAvail := minFunds

	ctl := map[address.Address]struct{}{}
	for _, a := range append(mi.ControlAddresses, mi.Owner, mi.Worker) {
		ctl[a] = struct{}{}
	}

	for _, addr := range addrs {
		if addr.Protocol() != address.ID {
			var err error
			addr, err = a.StateLookupID(ctx, addr, types.EmptyTSK)
			if err != nil {
				log.Warnw("looking up control address", "address", addr, "error", err)
				continue
			}
		}

		if _, ok := ctl[addr]; !ok {
			log.Warnw("non-control address configured for sending messages", "address", addr)
			continue
		}

		if maybeUseAddress(ctx, a, addr, goodFunds, &leastBad, &bestAvail) {
			return leastBad, bestAvail, nil
		}
	}

	log.Warnw("No address had enough funds to for full message Fee, selecting least bad address", "address", leastBad, "balance", types.FIL(bestAvail), "optimalFunds", types.FIL(goodFunds), "minFunds", types.FIL(minFunds))

	return leastBad, bestAvail, nil
}

func maybeUseAddress(ctx context.Context, a addrSelectApi, addr address.Address, goodFunds abi.TokenAmount, leastBad *address.Address, bestAvail *abi.TokenAmount) bool {
	b, err := a.WalletBalance(ctx, addr)
	if err != nil {
		log.Errorw("checking control address balance", "addr", addr, "error", err)
		return false
	}

	if b.GreaterThanEqual(goodFunds) {
		k, err := a.StateAccountKey(ctx, addr, types.EmptyTSK)
		if err != nil {
			log.Errorw("getting account key", "error", err)
			return false
		}

		have, err := a.WalletHas(ctx, k)
		if err != nil {
			log.Errorw("failed to check control address", "addr", addr, "error", err)
			return false
		}

		if !have {
			log.Errorw("don't have key", "key", k, "address", addr)
			return false
		}

		*leastBad = addr
		*bestAvail = b
		return true
	}

	if b.GreaterThan(*bestAvail) {
		*leastBad = addr
		*bestAvail = b
	}

	log.Warnw("address didn't have enough funds to send message", "address", addr, "required", types.FIL(goodFunds), "balance", types.FIL(b))
	return false
}