package control

import (
	"bytes"

	"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/commonflags"
	"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/key"
	object "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/modules/object"
	commonCmd "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/internal/common"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/services/control"
	rawclient "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/rpc/client"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id"

	"github.com/mr-tron/base58"
	"github.com/spf13/cobra"
)

const (
	FullInfoFlag      = "full"
	FullInfoFlagUsage = "Print full ShardInfo."
)

var locateObjectCmd = &cobra.Command{
	Use:   "locate-object",
	Short: "List shards storing the object",
	Long:  "List shards storing the object",
	Run:   locateObject,
}

func initControlLocateObjectCmd() {
	initControlFlags(locateObjectCmd)

	flags := locateObjectCmd.Flags()

	flags.String(commonflags.CIDFlag, "", commonflags.CIDFlagUsage)
	_ = locateObjectCmd.MarkFlagRequired(commonflags.CIDFlag)

	flags.String(commonflags.OIDFlag, "", commonflags.OIDFlagUsage)
	_ = locateObjectCmd.MarkFlagRequired(commonflags.OIDFlag)

	flags.Bool(commonflags.JSON, false, "Print shard info as a JSON array. Requires --full flag.")
	flags.Bool(FullInfoFlag, false, FullInfoFlagUsage)
}

func locateObject(cmd *cobra.Command, _ []string) {
	var cnr cid.ID
	var obj oid.ID

	_ = object.ReadObjectAddress(cmd, &cnr, &obj)

	pk := key.Get(cmd)

	body := new(control.ListShardsForObjectRequest_Body)
	body.SetContainerId(cnr.EncodeToString())
	body.SetObjectId(obj.EncodeToString())
	req := new(control.ListShardsForObjectRequest)
	req.SetBody(body)
	signRequest(cmd, pk, req)

	cli := getClient(cmd, pk)

	var err error
	var resp *control.ListShardsForObjectResponse
	err = cli.ExecRaw(func(client *rawclient.Client) error {
		resp, err = control.ListShardsForObject(client, req)
		return err
	})
	commonCmd.ExitOnErr(cmd, "rpc error: %w", err)

	verifyResponse(cmd, resp.GetSignature(), resp.GetBody())

	shardIDs := resp.GetBody().GetShard_ID()

	isFull, _ := cmd.Flags().GetBool(FullInfoFlag)
	if !isFull {
		for _, id := range shardIDs {
			cmd.Println(base58.Encode(id))
		}
		return
	}

	// get full shard info
	listShardsReq := new(control.ListShardsRequest)
	listShardsReq.SetBody(new(control.ListShardsRequest_Body))
	signRequest(cmd, pk, listShardsReq)
	var listShardsResp *control.ListShardsResponse
	err = cli.ExecRaw(func(client *rawclient.Client) error {
		listShardsResp, err = control.ListShards(client, listShardsReq)
		return err
	})
	commonCmd.ExitOnErr(cmd, "rpc error: %w", err)

	verifyResponse(cmd, listShardsResp.GetSignature(), listShardsResp.GetBody())

	shards := listShardsResp.GetBody().GetShards()
	sortShardsByID(shards)
	shards = filterShards(shards, shardIDs)

	isJSON, _ := cmd.Flags().GetBool(commonflags.JSON)
	if isJSON {
		prettyPrintShardsJSON(cmd, shards)
	} else {
		prettyPrintShards(cmd, shards)
	}
}

func filterShards(info []control.ShardInfo, ids [][]byte) []control.ShardInfo {
	var res []control.ShardInfo
	for _, id := range ids {
		for _, inf := range info {
			if bytes.Equal(inf.Shard_ID, id) {
				res = append(res, inf)
			}
		}
	}
	return res
}