package agent_execution_test

import (
	"context"
	"strconv"
	"testing"

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

type graphStore struct {
	graphs  map[string]work.WorkGraph
	queued  []string
	intents []string
}

func (s *graphStore) SaveWorkGraph(_ context.Context, graph work.WorkGraph) error {
	if s.graphs == nil {
		s.graphs = map[string]work.WorkGraph{}
	}
	s.graphs[graph.ID+":"+itoa(graph.Revision)] = graph
	return nil
}
func (s *graphStore) LoadWorkGraph(_ context.Context, id string, revision int64) (work.WorkGraph, error) {
	return s.graphs[id+":"+itoa(revision)], nil
}
func (s *graphStore) SaveDispatchIntent(_ context.Context, intent work.AgentDispatchIntent) error {
	s.intents = append(s.intents, intent.DispatchID)
	return nil
}
func (s *graphStore) EnqueueDispatch(_ context.Context, item work.AgentDispatchQueueItem) error {
	s.queued = append(s.queued, item.DispatchID)
	return nil
}

func itoa(value int64) string {
	return strconv.FormatInt(value, 10)
}

func TestGraphRuntimeSchedulesFanOutAndCompletesFanIn(t *testing.T) {
	store := &graphStore{}
	graph, ready, err := application.AdmitAndSchedule(context.Background(), store, application.GraphScheduleRequest{Graph: work.WorkGraph{ID: "g", Revision: 1, Budget: work.GraphBudget{MaxParallel: 2, MaxAttempts: 1}, Nodes: []work.GraphNode{{ID: "a", DispatchID: "d-a", IdempotencyKey: "a", State: work.GraphNodeReady}, {ID: "b", DispatchID: "d-b", IdempotencyKey: "b", State: work.GraphNodeReady}, {ID: "join", DispatchID: "d-join", IdempotencyKey: "join", State: work.GraphNodeBlocked}}, Edges: []work.GraphEdge{{From: "a", To: "join"}, {From: "b", To: "join"}}}, Intents: []work.AgentDispatchIntent{{DispatchID: "d-a"}, {DispatchID: "d-b"}, {DispatchID: "d-join"}}, Now: 1})
	if err != nil || len(ready) != 2 || len(store.queued) != 2 {
		t.Fatalf("graph=%+v ready=%v queued=%v err=%v", graph, ready, store.queued, err)
	}
	if _, err := application.CompleteGraphNode(context.Background(), store, "g", 2, "a", true); err != nil {
		t.Fatal(err)
	}
	if _, err := application.CompleteGraphNode(context.Background(), store, "g", 3, "b", true); err != nil {
		t.Fatal(err)
	}
	if len(store.queued) != 3 || store.queued[2] != "d-join" {
		t.Fatalf("fan-in dispatch was not enqueued: %v", store.queued)
	}
	completed, err := application.CompleteGraphNode(context.Background(), store, "g", 4, "join", true)
	if err != nil || completed.State != work.GraphCompleted {
		t.Fatalf("completed=%+v err=%v", completed, err)
	}
}