package worker

import (
	"context"
	"net"
	"testing"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/domain/workerruntime"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/grpc/test/bufconn"
	"google.golang.org/protobuf/types/known/timestamppb"
)

func TestGRPCHandlerMapsWorkerRequests(t *testing.T) {
	service := &grpcServiceStub{registration: Registration{WorkerID: "worker-a", SessionID: "session-a", OwnershipEpoch: 9, HeartbeatInterval: time.Second, MaxStaleAge: 2 * time.Second}}
	client, cleanup := grpcWorkerClient(t, service, allowIdentity{})
	defer cleanup()
	registered, err := client.RegisterWorker(context.Background(), &controlplanev1.RegisterWorkerRequest{
		WorkerId: "worker-a", InstanceId: "instance-a", Zone: "zone-a", SupportedProtocolVersion: 1,
	})
	if err != nil || registered.GetSessionId() != "session-a" || registered.GetHeartbeatInterval().AsDuration() != time.Second {
		t.Fatalf("RegisterWorker() = %+v, %v", registered, err)
	}
	if err := service.acknowledgeErr; err != nil {
		t.Fatal(err)
	}
	_, err = client.AcknowledgeSnapshot(context.Background(), &controlplanev1.AcknowledgeSnapshotRequest{
		WorkerId: "worker-a", SessionId: "session-a", Version: 7, OwnershipEpoch: 9, Checksum: make([]byte, 32),
	})
	if err != nil || service.acknowledgement.Version != 7 {
		t.Fatalf("AcknowledgeSnapshot() error = %v; command=%+v", err, service.acknowledgement)
	}
	response, err := client.ReportRuntime(context.Background(), &controlplanev1.ReportRuntimeRequest{
		WorkerId: "worker-a", SessionId: "session-a", SnapshotVersion: 7, OwnershipEpoch: 9, ReportSequence: 1,
		ObservedAt: timestamppb.New(time.Now()), Counters: []*controlplanev1.ProxyRuntime{{ProxyId: "proxy-a", Active: 2, Reserved: 1}},
	})
	if err != nil || response.GetAcceptedOwnershipEpoch() != 9 || service.report.Counters[0].Active != 2 {
		t.Fatalf("ReportRuntime() = %+v, %v; report=%+v", response, err, service.report)
	}
}

func TestGRPCHandlerMapsErrorsAndLeavesStreamsUnimplemented(t *testing.T) {
	service := &grpcServiceStub{registerErr: ErrProtocolVersion}
	client, cleanup := grpcWorkerClient(t, service, allowIdentity{})
	defer cleanup()
	_, err := client.RegisterWorker(context.Background(), &controlplanev1.RegisterWorkerRequest{WorkerId: "worker-a"})
	if status.Code(err) != codes.FailedPrecondition {
		t.Fatalf("RegisterWorker() code = %s, want FailedPrecondition", status.Code(err))
	}
	stream, err := client.WatchSnapshots(context.Background(), &controlplanev1.WatchSnapshotsRequest{})
	_, streamErr := stream.Recv()
	if err != nil || status.Code(streamErr) != codes.Unimplemented {
		t.Fatalf("WatchSnapshots() = %v, %v", err, streamErr)
	}
	outcomes, err := client.ReportOutcomes(context.Background())
	_, outcomesErr := outcomes.CloseAndRecv()
	if err != nil || status.Code(outcomesErr) != codes.Unimplemented {
		t.Fatalf("ReportOutcomes() = %v, %v", err, outcomesErr)
	}
}

type grpcServiceStub struct {
	registration    Registration
	registerErr     error
	acknowledgement SnapshotAcknowledgement
	acknowledgeErr  error
	report          workerruntime.Report
	reportErr       error
}

func (stub *grpcServiceStub) Register(context.Context, RegisterCommand) (Registration, error) {
	return stub.registration, stub.registerErr
}
func (stub *grpcServiceStub) Acknowledge(_ context.Context, acknowledgement SnapshotAcknowledgement) error {
	stub.acknowledgement = acknowledgement
	return stub.acknowledgeErr
}
func (stub *grpcServiceStub) ReportRuntime(_ context.Context, report workerruntime.Report) (RuntimeDecision, error) {
	stub.report = report
	return RuntimeDecision{AcceptedOwnershipEpoch: 9}, stub.reportErr
}

type allowIdentity struct{}

func (allowIdentity) Authorize(context.Context, string) error { return nil }

func grpcWorkerClient(t *testing.T, service Service, identity IdentityAuthorizer) (controlplanev1.WorkerControlPlaneClient, func()) {
	t.Helper()
	listener := bufconn.Listen(1 << 20)
	server := grpc.NewServer()
	controlplanev1.RegisterWorkerControlPlaneServer(server, NewGRPCHandler(service, identity))
	go func() { _ = server.Serve(listener) }()
	connection, err := grpc.NewClient("passthrough:///bufnet", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), grpc.WithInsecure())
	if err != nil {
		t.Fatalf("grpc.NewClient(): %v", err)
	}
	return controlplanev1.NewWorkerControlPlaneClient(connection), func() { _ = connection.Close(); server.Stop(); _ = listener.Close() }
}