package controlplane

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/gateway/snapshot"

	"google.golang.org/protobuf/types/known/durationpb"
	"google.golang.org/protobuf/types/known/emptypb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

func TestSessionRunnerRegistersOnceWatchesAndReports(t *testing.T) {
	store := snapshot.NewStore("cluster-a", "worker-a")
	full := &controlplanev1.WorkerSnapshot{
		Version: 1, OwnershipEpoch: 7, GeneratedAt: timestamppb.New(time.Now()), ValidUntil: timestamppb.New(time.Now().Add(time.Minute)),
	}
	setSnapshotChecksum(t, full)
	runtime := make(chan *controlplanev1.ReportRuntimeRequest, 8)
	reporterClient := &sessionReporterClient{runtime: runtime}
	reporter, err := NewRuntimeReporter(reporterClient, store, Options{
		WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1, Now: time.Now,
	})
	if err != nil {
		t.Fatalf("NewRuntimeReporter(): %v", err)
	}
	watcherClient := &sessionWatcherClient{full: full}
	watcher, err := NewSnapshotWatcher(watcherClient, store, SnapshotWatcherOptions{ClusterID: "cluster-a", WorkerID: "worker-a"})
	if err != nil {
		t.Fatalf("NewSnapshotWatcher(): %v", err)
	}
	runner, err := NewSessionRunner(reporter, watcher)
	if err != nil {
		t.Fatalf("NewSessionRunner(): %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	result := make(chan error, 1)
	go func() { result <- runner.Run(ctx) }()
	select {
	case request := <-runtime:
		if request.GetSessionId() != "session-a" || request.GetSnapshotVersion() != 1 || request.GetReportSequence() != 1 {
			t.Fatalf("runtime request = %+v", request)
		}
	case <-time.After(time.Second):
		t.Fatal("runtime report was not sent")
	}
	cancel()
	if err := <-result; !errors.Is(err, context.Canceled) {
		t.Fatalf("Run() error = %v, want context canceled", err)
	}
	if reporterClient.registers != 1 || watcherClient.ack == nil || !watcherClient.ack.GetApplied() {
		t.Fatalf("registers=%d ack=%+v", reporterClient.registers, watcherClient.ack)
	}
}

type sessionReporterClient struct {
	mu        sync.Mutex
	registers int
	runtime   chan *controlplanev1.ReportRuntimeRequest
}

func (client *sessionReporterClient) RegisterWorker(context.Context, *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) {
	client.mu.Lock()
	client.registers++
	client.mu.Unlock()
	return &controlplanev1.RegisterWorkerResponse{
		WorkerId: "worker-a", SessionId: "session-a", OwnershipEpoch: 7,
		HeartbeatInterval: durationpb.New(5 * time.Millisecond), MaxStaleAge: durationpb.New(time.Second),
	}, nil
}

func (client *sessionReporterClient) ReportRuntime(_ context.Context, request *controlplanev1.ReportRuntimeRequest) (*controlplanev1.ReportRuntimeResponse, error) {
	client.runtime <- request
	return &controlplanev1.ReportRuntimeResponse{AcceptedOwnershipEpoch: 7}, nil
}

type sessionWatcherClient struct {
	full *controlplanev1.WorkerSnapshot
	ack  *controlplanev1.AcknowledgeSnapshotRequest
}

func (client *sessionWatcherClient) Watch(ctx context.Context, _ *controlplanev1.WatchSnapshotsRequest) (SnapshotStream, error) {
	return &sessionSnapshotStream{ctx: ctx, full: client.full}, nil
}

func (client *sessionWatcherClient) Acknowledge(_ context.Context, acknowledgement *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) {
	client.ack = acknowledgement
	return &emptypb.Empty{}, nil
}

type sessionSnapshotStream struct {
	ctx  context.Context
	full *controlplanev1.WorkerSnapshot
	sent bool
}

func (stream *sessionSnapshotStream) Recv() (*controlplanev1.SnapshotEnvelope, error) {
	if !stream.sent {
		stream.sent = true
		return &controlplanev1.SnapshotEnvelope{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: stream.full}}, nil
	}
	<-stream.ctx.Done()
	return nil, stream.ctx.Err()
}