package orchestrator_test

import (
	"context"
	"path/filepath"
	"testing"
	"time"

	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/envelope"
	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/event"
	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/orchestrator"
	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/org"
	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/prompt"
	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/runtime"
	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/runtime/scripted"
	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/store"
	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/transport"
)

// TestPhaseN_OrgRunnerSpawnsFromYAML exercises the end-to-end Phase N flow:
// build an org Definition in-memory, OrgRunner spawns every role as an
// agent, Master delegates via the loaded delegates_to graph (not hardcoded
// "worker-1"), and the delegate reaches the right agent id.
func TestPhaseN_OrgRunnerSpawnsFromYAML(t *testing.T) {
	tmp := t.TempDir()
	st, err := store.Open(filepath.Join(tmp, "harness.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer st.Close()
	bus := event.NewBus(st)
	q := transport.New(st, bus)
	orch := orchestrator.New(st, bus)
	rt := scripted.New()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// In-memory org so we don't depend on disk yaml in this test.
	def := &org.Definition{
		Name: "tiny-hier", Version: 1,
		SubOrgs: []org.SubOrg{{ID: "frontend", Teams: []string{"fe-team"}, ZoneScope: []string{"client/**/*"}}},
		Roles: []org.Role{
			{ID: "master", Provider: "scripted", DelegatesTo: []string{"fe-lead"}, Tools: []string{"delegate"}},
			{ID: "fe-lead", Provider: "scripted", DelegatesTo: []string{"fe-worker"}, Tools: []string{"send_message"}},
			{ID: "fe-worker", Provider: "scripted", Tools: []string{"introspect", "write_file", "send_message"}},
		},
		Teams: []org.Team{{ID: "fe-team", Roster: []string{"fe-lead", "fe-worker"}}},
	}

	runtimes := map[string]runtime.Runtime{
		"default":  rt,
		"scripted": rt,
	}
	runner := orchestrator.NewOrgRunner(orch, q, runtimes, def)

	_ = orch.CreateRun(ctx, "run-n", "phase N")
	spawned, err := runner.SpawnAll(ctx, "run-n")
	if err != nil {
		t.Fatalf("SpawnAll: %v", err)
	}
	if len(spawned) != 3 {
		t.Errorf("spawned = %d, want 3", len(spawned))
	}

	masterID, _ := runner.AgentIDFor("master")
	leadID, _ := runner.AgentIDFor("fe-lead")
	workerID, _ := runner.AgentIDFor("fe-worker")
	if masterID == "" || leadID == "" || workerID == "" {
		t.Fatalf("agent ids: master=%s lead=%s worker=%s", masterID, leadID, workerID)
	}

	target, ok := runner.DelegateTarget("master")
	if !ok || target != "fe-lead" {
		t.Errorf("delegate target = %s (ok=%v), want fe-lead", target, ok)
	}

	// Master uses the runner now (no hardcoded worker-1).
	m := orchestrator.NewMaster(orch, q, prompt.NewRules(), "run-n")
	m.Runner = runner

	// Script the lead to forward via send_message to the worker.
	rt.SetScript(leadID, []*runtime.LLMResponse{
		{
			Text:   "forwarding to worker",
			Tokens: runtime.TokenUsage{Prompt: 50, Completion: 25},
			ToolCalls: []runtime.ToolCall{
				{Name: "send_message", Args: map[string]any{
					"to":      workerID,
					"type":    "delegate",
					"ttl_ms":  60000,
					"id":      "msg_lead_to_worker",
					"run_id":  "run-n",
					"payload": map[string]any{"intent": "implement it", "expects": "report"},
				}},
			},
		},
	})

	if _, _, err := determDispatch(t, m, "Add a tooltip to the dashboard"); err != nil {
		t.Fatalf("Dispatch: %v", err)
	}

	// Lead inbox should have the delegate.
	leadInbox, err := q.Receive(ctx, leadID)
	if err != nil {
		t.Fatal(err)
	}
	if leadInbox == nil {
		t.Fatal("lead never received the delegate from master")
	}
	if leadInbox.From != masterID {
		t.Errorf("delegate from = %s, want %s", leadInbox.From, masterID)
	}
	// Ack and let the lead's scripted turn fire via HandleOne against a
	// runtime.Agent. We construct a fresh agent for the lead here so we can
	// drive one turn.
	_ = q.Ack(ctx, leadInbox.ID)
	leadAgent := runtime.NewAgent(rt, st, bus, q, leadID, "fe-lead", "run-n",
		nil, []string{"send_message"}, "scripted", 0.005)
	// Re-send the message so HandleOne can consume it: just send a new one.
	again := &envelope.Envelope{
		ID: "msg_again", RunID: "run-n", From: masterID, To: leadID,
		Type: envelope.TypeDelegate, TTLMs: 60000,
		Payload: envelope.Payload{Intent: "do it", Expects: envelope.ExpectsReport},
	}
	if err := q.Send(ctx, again); err != nil {
		t.Fatal(err)
	}
	if h, err := leadAgent.HandleOne(ctx); err != nil || !h {
		t.Fatalf("leadAgent.HandleOne: handled=%v err=%v", h, err)
	}

	got, err := q.Receive(ctx, workerID)
	if err != nil {
		t.Fatal(err)
	}
	if got == nil || got.Type != envelope.TypeDelegate {
		t.Fatalf("worker delegate missing or wrong type: %+v", got)
	}
	if got.From != leadID {
		t.Errorf("worker delegate from = %s, want %s", got.From, leadID)
	}
}
