package sourceagent

import (
	"bytes"
	"context"
	"testing"
	"time"

	"github.com/local/glpi-neural-brain/internal/vectorgraph"
)

func TestVectorGraphComputeWireRoundTrip(t *testing.T) {
	request := VectorGraphComputeRequest{
		Header: VectorGraphComputeHeader{
			SchemaVersion:  SchemaVersion,
			JobID:          "compute-test",
			Kind:           ComputeKindVectorGraph,
			GraphVersion:   42,
			Primary:        vectorgraph.Config{Neighbors: 2, CandidateLimit: 8, HashBits: 8, HashTables: 2, BandBits: 4, MinSimilarity: .8, MinAffinity: .2},
			OrphanPass:     true,
			OrphanFocusIDs: []string{"c"},
		},
		Entries: []vectorgraph.Entry{
			{ID: "a", Vector: []float32{1, 0, 0}},
			{ID: "b", Vector: []float32{.99, .01, 0}},
			{ID: "c", Vector: []float32{.98, .02, 0}},
		},
	}
	var buffer bytes.Buffer
	if err := WriteVectorGraphJob(&buffer, request); err != nil {
		t.Fatal(err)
	}
	decoded, err := ReadVectorGraphJob(bytes.NewReader(buffer.Bytes()), 1<<20)
	if err != nil {
		t.Fatal(err)
	}
	if decoded.Header.JobID != request.Header.JobID || decoded.Header.GraphVersion != 42 || len(decoded.Entries) != 3 {
		t.Fatalf("unexpected decoded job: %+v", decoded.Header)
	}
	for i := range request.Entries {
		if decoded.Entries[i].ID != request.Entries[i].ID || len(decoded.Entries[i].Vector) != len(request.Entries[i].Vector) {
			t.Fatalf("entry %d mismatch: got=%+v want=%+v", i, decoded.Entries[i], request.Entries[i])
		}
		for j := range request.Entries[i].Vector {
			if decoded.Entries[i].Vector[j] != request.Entries[i].Vector[j] {
				t.Fatalf("entry %d vector %d mismatch", i, j)
			}
		}
	}
}

func TestExecuteVectorGraphJobRunsOrphanSecondPassWithoutModel(t *testing.T) {
	entries := []vectorgraph.Entry{
		{ID: "a", Vector: []float32{1, 0, 0}},
		{ID: "b", Vector: []float32{.99, .01, 0}},
		{ID: "orphan", Vector: []float32{.97, .03, 0}},
		{ID: "far", Vector: []float32{0, 1, 0}},
	}
	request := VectorGraphComputeRequest{Header: VectorGraphComputeHeader{
		SchemaVersion: SchemaVersion, JobID: "compute-orphan", Kind: ComputeKindVectorGraph, GraphVersion: 7,
		Primary:        vectorgraph.Config{Neighbors: 1, CandidateLimit: 4, HashBits: 8, HashTables: 2, BandBits: 4, MinSimilarity: .995, MinAffinity: .2},
		OrphanPass:     true,
		Orphan:         vectorgraph.Config{Neighbors: 1, CandidateLimit: 8, HashBits: 8, HashTables: 2, BandBits: 4, MinSimilarity: .90, MinAffinity: .1},
		OrphanFocusIDs: []string{"orphan"},
	}, Entries: entries}
	result := ExecuteVectorGraphJob(request)
	if result.Kind != ComputeKindVectorGraph || result.GraphVersion != 7 || result.Error != "" {
		t.Fatalf("unexpected result: %+v", result)
	}
	if result.Orphan.Stats.Focused == 0 {
		t.Fatalf("orphan pass did not inspect requested focus: %+v", result.Orphan.Stats)
	}
	for _, link := range result.Orphan.Links {
		if link.Source != "orphan" && link.Target != "orphan" {
			t.Fatalf("orphan pass emitted unrelated link: %+v", link)
		}
	}
}

func TestVectorGraphComputeBrokerClaimComplete(t *testing.T) {
	store := &Store{computeJobs: map[string]*computeJobState{}}
	request := VectorGraphComputeRequest{Header: VectorGraphComputeHeader{JobID: "broker-job", Kind: ComputeKindVectorGraph, GraphVersion: 9}, Entries: []vectorgraph.Entry{{ID: "a", Vector: []float32{1, 0}}, {ID: "b", Vector: []float32{.99, .01}}}}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	resultCh := make(chan VectorGraphComputeResult, 1)
	errCh := make(chan error, 1)
	go func() {
		result, err := store.SubmitVectorGraphJob(ctx, request)
		if err != nil {
			errCh <- err
			return
		}
		resultCh <- result
	}()

	deadline := time.Now().Add(time.Second)
	var claimed VectorGraphComputeRequest
	var ok bool
	for time.Now().Before(deadline) {
		claimed, ok = store.ClaimVectorGraphJob("agent-cpu", time.Minute)
		if ok {
			break
		}
		time.Sleep(time.Millisecond)
	}
	if !ok || claimed.Header.JobID != "broker-job" {
		t.Fatalf("compute job was not claimable: ok=%v job=%+v", ok, claimed.Header)
	}
	result := ExecuteVectorGraphJob(claimed)
	if err := store.CompleteVectorGraphJob("agent-cpu", result); err != nil {
		t.Fatal(err)
	}
	select {
	case err := <-errCh:
		t.Fatal(err)
	case got := <-resultCh:
		if got.AgentID != "agent-cpu" || got.GraphVersion != 9 || got.Kind != ComputeKindVectorGraph {
			t.Fatalf("unexpected completed result: %+v", got)
		}
	case <-ctx.Done():
		t.Fatal(ctx.Err())
	}
}