package application

import (
	"context"
	"testing"
	"time"

	"git.cin.su/m/Assistant/contracts/work"
)

type bridgeStore struct {
	callbacks []work.AgentCallback
	lease     work.AgentLease
}

func (s *bridgeStore) SaveDispatchIntent(context.Context, work.AgentDispatchIntent) error { return nil }
func (s *bridgeStore) SaveAgentLease(_ context.Context, lease work.AgentLease) error {
	s.lease = lease
	return nil
}
func (s *bridgeStore) RenewAgentLease(_ context.Context, lease work.AgentLease, _ time.Time) error {
	s.lease = lease
	return nil
}
func (s *bridgeStore) AcceptAgentCallback(_ context.Context, callback work.AgentCallback) error {
	s.callbacks = append(s.callbacks, callback)
	return nil
}

type bridgeExecutor struct{}

func (bridgeExecutor) Execute(context.Context, work.AgentExecutionRequest) (work.AgentExecutionResult, error) {
	return work.AgentExecutionResult{State: work.AgentExecutionCompleted, Summary: "done", ObservedAtUnixMillis: 2}, nil
}

type slowBridgeExecutor struct{}

func (slowBridgeExecutor) Execute(ctx context.Context, request work.AgentExecutionRequest) (work.AgentExecutionResult, error) {
	select {
	case <-time.After(45 * time.Millisecond):
		return work.AgentExecutionResult{State: work.AgentExecutionCompleted, Summary: "done", ObservedAtUnixMillis: time.Now().UnixMilli()}, nil
	case <-ctx.Done():
		return work.AgentExecutionResult{}, ctx.Err()
	}
}

func TestExecuteAgentDispatchPersistsLeaseAndBridgeCallbacks(t *testing.T) {
	store := &bridgeStore{}
	now := time.UnixMilli(1)
	result, err := ExecuteAgentDispatch(context.Background(), store, bridgeExecutor{}, work.AgentDispatchIntent{DispatchID: "d", IdempotencyKey: "i", DelegationID: "del", RunID: "run", ImplementationID: "impl", ImplementationRevision: 1, SourceRevision: "source", WorktreeID: "tree", VerificationPlanRef: "verify", CreatedAtUnixMillis: 1}, "worker", 1, now, time.Minute)
	if err != nil || result.State != work.AgentExecutionCompleted {
		t.Fatalf("result=%+v err=%v", result, err)
	}
	if len(store.callbacks) != 2 || store.callbacks[0].Kind != work.AgentCallbackHeartbeat || store.callbacks[1].Kind != work.AgentCallbackResult {
		t.Fatalf("callbacks=%+v", store.callbacks)
	}
}

func TestExecuteAgentDispatchRenewsLeaseAndOrdersHeartbeats(t *testing.T) {
	store := &bridgeStore{}
	now := time.Now().UTC()
	_, err := ExecuteAgentDispatch(context.Background(), store, slowBridgeExecutor{}, work.AgentDispatchIntent{DispatchID: "d-slow", IdempotencyKey: "i-slow", DelegationID: "del-slow", RunID: "run-slow", ImplementationID: "impl", ImplementationRevision: 1, SourceRevision: "source", WorktreeID: "tree", VerificationPlanRef: "verify", CreatedAtUnixMillis: now.UnixMilli()}, "worker", 1, now, 30*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	if len(store.callbacks) < 3 {
		t.Fatalf("expected heartbeat callbacks, got %d", len(store.callbacks))
	}
	for i, callback := range store.callbacks {
		if callback.Sequence != uint64(i+1) {
			t.Fatalf("callback sequence is not ordered: %#v", store.callbacks)
		}
	}
}