package worker

import (
	"context"
	"errors"
	"time"

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

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

var ErrSnapshotsUnavailable = errors.New("worker snapshots are unavailable")

type SnapshotWatchRequest struct {
	WorkerID           string
	SessionID          string
	LastAppliedVersion uint64
	LastChecksum       []byte
}

type SnapshotSource interface {
	Watch(context.Context, SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error)
}

type OwnershipEpochReader interface {
	CurrentOwnershipEpoch(context.Context) (uint64, error)
}

type InitialSnapshotSource struct {
	epochs   OwnershipEpochReader
	validFor time.Duration
	now      func() time.Time
}

func NewInitialSnapshotSource(epochs OwnershipEpochReader, validFor time.Duration, now func() time.Time) (*InitialSnapshotSource, error) {
	if epochs == nil || validFor <= 0 || now == nil {
		return nil, ErrSnapshotsUnavailable
	}
	return &InitialSnapshotSource{epochs: epochs, validFor: validFor, now: now}, nil
}

func (source *InitialSnapshotSource) Watch(ctx context.Context, request SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) {
	if source == nil || ctx == nil || !workerruntime.ValidIdentifier(request.WorkerID) ||
		!workerruntime.ValidIdentifier(request.SessionID) || request.LastAppliedVersion == ^uint64(0) {
		return nil, ErrSnapshotsUnavailable
	}
	epoch, err := source.epochs.CurrentOwnershipEpoch(ctx)
	if err != nil {
		return nil, err
	}
	if epoch == 0 {
		return nil, ErrSnapshotsUnavailable
	}
	now := source.now().UTC()
	if now.IsZero() {
		return nil, ErrSnapshotsUnavailable
	}
	full := &controlplanev1.WorkerSnapshot{
		Version: request.LastAppliedVersion + 1, OwnershipEpoch: epoch,
		GeneratedAt: timestamppb.New(now), ValidUntil: timestamppb.New(now.Add(source.validFor)),
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil {
		return nil, err
	}
	full.Checksum = append([]byte(nil), checksum[:]...)
	updates := make(chan *controlplanev1.WorkerSnapshot, 1)
	updates <- full
	return updates, nil
}

var _ SnapshotSource = (*InitialSnapshotSource)(nil)