package router

import (
	"encoding/json"
	"testing"
	"time"

	"orchestra/internal/authz"
	"orchestra/internal/domain"
	"orchestra/internal/registry"
	"orchestra/internal/store"
)

func TestAssignPendingUsesProjectAffinityAndCapability(t *testing.T) {
	s, err := store.Open(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	add := func(id, project string, caps []string) {
		b, _ := json.Marshal(map[string]any{"source": "test", "external_id": id, "project": project, "capability": caps})
		if err := s.Append(domain.Event{Type: "TaskCreated", TaskID: id, Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
			t.Fatal(err)
		}
	}
	add("task", "mechanical", []string{"mechanical"})
	r, err := registry.New(registry.Config{
		Projects: []registry.Project{{ID: "mechanical", MachineAffinity: []string{"homesrv"}}},
		Machines: []registry.Machine{{ID: "homesrv", Address: "homesrv:1"}, {ID: "workpc", Address: "workpc:1"}},
		Herdrs:   []registry.Herdr{{ID: "wrong-machine", MachineID: "workpc", Capabilities: []string{"mechanical"}}, {ID: "wrong-cap", MachineID: "homesrv", Capabilities: []string{"correlation"}}},
	})
	if err != nil {
		t.Fatal(err)
	}
	out, err := (&Router{Store: s, Registry: r}).AssignPending()
	if err != nil {
		t.Fatal(err)
	}
	if len(out) != 0 {
		t.Fatalf("assigned despite affinity/capability mismatch: %#v", out)
	}
}

func TestAssignPendingDoesNotPreemptRunningWork(t *testing.T) {
	s, err := store.Open(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	b, _ := json.Marshal(map[string]any{"source": "test", "external_id": "queued", "project": "p", "capability": []string{}})
	if err := s.Append(domain.Event{Type: "TaskCreated", TaskID: "queued", Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
		t.Fatal(err)
	}
	r, err := registry.New(registry.Config{Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}}, Machines: []registry.Machine{{ID: "m", Address: "m:1"}}, Herdrs: []registry.Herdr{{ID: "h", MachineID: "m", Concurrency: 1}}})
	if err != nil {
		t.Fatal(err)
	}
	_, err = s.Lease("queued", "h", time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := (&Router{Store: s, Registry: r}).AssignPending(); err != nil {
		t.Fatal(err)
	}
	task, ok := s.Task("queued")
	if !ok || task.State != domain.StateLeased || task.Lease == nil || task.Lease.HarnessID != "h" {
		t.Fatalf("running lease was changed while assigning: %#v", task)
	}
}