package weed_server

import (
	"context"
	"testing"
	"time"

	"github.com/seaweedfs/raft"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
	"github.com/seaweedfs/seaweedfs/weed/sequence"
	"github.com/seaweedfs/seaweedfs/weed/storage/needle"
	"github.com/seaweedfs/seaweedfs/weed/storage/super_block"
	"github.com/seaweedfs/seaweedfs/weed/storage/types"
	"github.com/seaweedfs/seaweedfs/weed/topology"
)

// leaderRaftServer answers only State(); the embedded nil interface panics if any
// other method is called.
type leaderRaftServer struct {
	raft.Server
}

func (leaderRaftServer) State() string { return raft.Leader }

func newLeaderMaster() *MasterServer {
	topo := topology.NewTopology("test", sequence.NewMemorySequencer(), 1<<20, 5, false)
	topo.RaftServer = leaderRaftServer{}
	return &MasterServer{
		Topo:                    topo,
		option:                  &MasterOption{},
		volumeGrowthRequestChan: make(chan *topology.VolumeGrowRequest, 1<<6),
	}
}

// markGrowthInFlight flags growth on the same VolumeLayout the handler resolves,
// by mirroring its exact lookup.
func markGrowthInFlight(t *testing.T, topo *topology.Topology, req *master_pb.AssignRequest) {
	t.Helper()
	rp, err := super_block.NewReplicaPlacementFromString(req.Replication)
	require.NoError(t, err)
	ttl, err := needle.ReadTTL(req.Ttl)
	require.NoError(t, err)
	topo.GetVolumeLayout(req.Collection, rp, ttl, types.ToDiskType(req.DiskType)).AddGrowRequest()
}

// With free space but no writable volume and growth already in flight, Assign
// sheds with a retryable ResourceExhausted immediately instead of spinning to
// timeout. ResourceExhausted (not Unavailable) so the shed isn't mistaken for a
// dead channel and doesn't tear down the shared master connection.
func TestAssignShedsLoadWhenGrowthInFlight(t *testing.T) {
	ms := newLeaderMaster()
	// Free volume slots but no writable volume yet, so growth is the remedy.
	ms.Topo.GetOrCreateDataCenter("dc1").GetOrCreateRack("rack1").
		GetOrCreateDataNode("127.0.0.1", 8080, 18080, "127.0.0.1", "dn1", map[string]uint32{"": 100})

	req := &master_pb.AssignRequest{Count: 1, Replication: "000"}
	markGrowthInFlight(t, ms.Topo, req)

	start := time.Now()
	resp, err := ms.Assign(context.Background(), req)
	elapsed := time.Since(start)

	require.Error(t, err)
	require.Nil(t, resp)
	st, ok := status.FromError(err)
	require.True(t, ok)
	assert.Equal(t, codes.ResourceExhausted, st.Code())
	// Shed immediately rather than spinning out the 10s retry budget.
	assert.Less(t, elapsed, 2*time.Second)
	// Growth was already pending, so we must not enqueue another grow request.
	assert.Len(t, ms.volumeGrowthRequestChan, 0)
}

// Out of space, Assign fails fast with the real error rather than masking it as
// a retryable "growth in progress".
func TestAssignFailsFastWhenOutOfSpace(t *testing.T) {
	ms := newLeaderMaster() // no data nodes -> no free space

	req := &master_pb.AssignRequest{Count: 1, Replication: "000"}

	start := time.Now()
	resp, err := ms.Assign(context.Background(), req)
	elapsed := time.Since(start)

	require.Error(t, err)
	require.Nil(t, resp)
	if st, ok := status.FromError(err); ok {
		assert.NotEqual(t, codes.Unavailable, st.Code())
	}
	assert.Contains(t, err.Error(), "no free volumes left")
	assert.Less(t, elapsed, 2*time.Second)
}