package scheduler

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"testing"
	"time"

	"github.com/riverqueue/river"

	"github.com/nomadcode/nomadcode-core/internal/agent"
	"github.com/nomadcode/nomadcode-core/internal/authoring"
	"github.com/nomadcode/nomadcode-core/internal/db"
	"github.com/nomadcode/nomadcode-core/internal/model"
	"github.com/nomadcode/nomadcode-core/internal/notification"
	"github.com/nomadcode/nomadcode-core/internal/projectsync"
	"github.com/nomadcode/nomadcode-core/internal/roadmapsync"
	"github.com/nomadcode/nomadcode-core/internal/roadmapsyncpipeline"
	"github.com/nomadcode/nomadcode-core/internal/storage"
	"github.com/nomadcode/nomadcode-core/internal/workflow"
	"github.com/nomadcode/nomadcode-core/internal/workitem"
)

func TestRunAgentTaskCompletesFromMessage(t *testing.T) {
	worker := &TaskWorker{
		Agent: fakeAgentClient{
			result: agent.SendMessageResult{
				Message: &agent.Message{
					Kind:      "message",
					Role:      "agent",
					MessageID: "msg-1",
					Parts:     []agent.Part{{Kind: "text", Text: "done"}},
				},
				Raw: json.RawMessage(`{"kind":"message","messageId":"msg-1"}`),
			},
		},
	}

	raw, message, err := worker.runAgentTask(context.Background(), storage.Task{
		ID:      "task-1",
		Title:   "Fix issue",
		Source:  "manual",
		Payload: json.RawMessage(`{"prompt":"please fix it"}`),
	})
	if err != nil {
		t.Fatalf("runAgentTask returned error: %v", err)
	}
	if message != "done" {
		t.Fatalf("unexpected message: %q", message)
	}

	var output map[string]any
	if err := json.Unmarshal(raw, &output); err != nil {
		t.Fatalf("decode output: %v", err)
	}
	if output["mode"] != "a2a" || output["message"] != "done" || output["message_id"] != "msg-1" {
		t.Fatalf("unexpected output: %#v", output)
	}
}

func TestRunAgentTaskRejectsNonTerminalTask(t *testing.T) {
	worker := &TaskWorker{
		Agent: fakeAgentClient{
			result: agent.SendMessageResult{
				Task: &agent.Task{
					Kind:   "task",
					ID:     "remote-1",
					Status: agent.TaskStatus{State: agent.TaskStateWorking},
				},
			},
		},
	}

	_, _, err := worker.runAgentTask(context.Background(), storage.Task{
		ID:      "task-1",
		Title:   "Fix issue",
		Source:  "manual",
		Payload: json.RawMessage(`{}`),
	})
	if err == nil || !strings.Contains(err.Error(), "non-terminal") {
		t.Fatalf("expected non-terminal state error, got %v", err)
	}
}

func TestBuildAgentInputUsesPromptPayload(t *testing.T) {
	input := buildAgentInput(storage.Task{
		ID:      "task-1",
		Title:   "Fallback title",
		Source:  "manual",
		Payload: json.RawMessage(`{"prompt":"do the thing","instructions":"be concise"}`),
	})

	if !input.Blocking {
		t.Fatal("expected blocking A2A call")
	}
	if !strings.Contains(input.Text, "be concise") || !strings.Contains(input.Text, "do the thing") {
		t.Fatalf("unexpected text: %q", input.Text)
	}
	if input.Metadata["nomadcode_task_id"] != "task-1" {
		t.Fatalf("unexpected metadata: %#v", input.Metadata)
	}
}

func TestBuildGenerateInputIncludesTaskMetadata(t *testing.T) {
	input := buildGenerateInput(storage.Task{
		ID:      "task-1",
		Title:   "Fallback title",
		Source:  "plane",
		Payload: json.RawMessage(`{"prompt":"do the thing","instructions":"be concise"}`),
	})

	if input.Input != "do the thing" {
		t.Fatalf("unexpected input: %q", input.Input)
	}
	if input.Instructions != "be concise" {
		t.Fatalf("unexpected instructions: %q", input.Instructions)
	}
	if input.Metadata["task_id"] != "task-1" {
		t.Fatalf("unexpected metadata: %#v", input.Metadata)
	}
	if _, ok := input.Metadata["source"]; ok {
		t.Fatalf("metadata.source must not be sent to IOP-compatible responses: %#v", input.Metadata)
	}
}

type fakeAgentClient struct {
	result agent.SendMessageResult
	err    error
}

func (c fakeAgentClient) SendMessage(context.Context, agent.SendMessageInput) (agent.SendMessageResult, error) {
	return c.result, c.err
}

func (c fakeAgentClient) GetTask(context.Context, agent.GetTaskInput) (agent.Task, error) {
	return agent.Task{}, nil
}

func (c fakeAgentClient) CancelTask(context.Context, agent.CancelTaskInput) (agent.Task, error) {
	return agent.Task{}, nil
}

type fakeTaskLifecycle struct {
	started        []string
	completed      []string
	failed         []string
	failMessages   []string
	failInputs     []workflow.FailureInput
	mergedMetadata []map[string]any

	task storage.Task
}

type fakeWorkspaceSlotUpdater struct {
	calls []db.UpdateWorkspaceSlotStateParams
	err   error
}

func (f *fakeWorkspaceSlotUpdater) UpdateWorkspaceSlotState(_ context.Context, args db.UpdateWorkspaceSlotStateParams) (db.WorkspaceSlot, error) {
	f.calls = append(f.calls, args)
	if f.err != nil {
		return db.WorkspaceSlot{}, f.err
	}
	return db.WorkspaceSlot{
		ID:    args.ID,
		State: args.State,
	}, nil
}

func (f *fakeTaskLifecycle) StartTask(ctx context.Context, id string) (storage.Task, error) {
	if f.task.Status != "running" && f.task.Status != "pending" && f.task.Status != "queued" && f.task.Status != "failed" && f.task.Status != "" {
		return storage.Task{}, errors.New("invalid transition to running")
	}
	f.started = append(f.started, id)
	f.task.Status = "running"
	attempt := len(f.started)
	// Merge attempt into existing metadata rather than replacing it, so that
	// checkout metadata (used for authoring task detection) is preserved.
	var existing map[string]any
	if len(f.task.Metadata) > 0 && string(f.task.Metadata) != "null" {
		_ = json.Unmarshal(f.task.Metadata, &existing)
	}
	if existing == nil {
		existing = make(map[string]any)
	}
	existing["attempt"] = attempt
	f.task.Metadata, _ = json.Marshal(existing)
	f.task.UpdatedAt = time.Now().UTC()
	return f.task, nil
}

func (f *fakeTaskLifecycle) CompleteTask(ctx context.Context, id string, result json.RawMessage) (storage.Task, error) {
	if f.task.Status != "running" {
		return storage.Task{}, errors.New("invalid transition to completed")
	}
	f.completed = append(f.completed, id)
	f.task.Status = "completed"
	attempt := len(f.started)
	f.task.Metadata = json.RawMessage(fmt.Sprintf(`{"attempt":%d}`, attempt))
	f.task.UpdatedAt = time.Now().UTC()
	return f.task, nil
}

func (f *fakeTaskLifecycle) FailTaskWithMetadata(ctx context.Context, id string, input workflow.FailureInput) (storage.Task, error) {
	if f.task.Status != "running" {
		return storage.Task{}, errors.New("invalid transition to failed")
	}
	f.failed = append(f.failed, id)
	f.failMessages = append(f.failMessages, input.Message)
	f.failInputs = append(f.failInputs, input)
	f.task.Status = "failed"
	attempt := len(f.started)
	f.task.Metadata = json.RawMessage(fmt.Sprintf(`{"attempt":%d,"status_reason":"%s"}`, attempt, input.Message))
	f.task.UpdatedAt = time.Now().UTC()
	return f.task, nil
}

func (f *fakeTaskLifecycle) FailTask(ctx context.Context, id string, message string) (storage.Task, error) {
	return f.FailTaskWithMetadata(ctx, id, workflow.FailureInput{
		Message: message,
		Type:    workflow.FailureTypeExecution,
	})
}

func (f *fakeTaskLifecycle) MergeTaskMetadata(_ context.Context, _ string, updates map[string]any) (storage.Task, error) {
	if f.task.Status == "completed" || f.task.Status == "failed" {
		return f.task, nil
	}
	f.mergedMetadata = append(f.mergedMetadata, updates)
	var current map[string]any
	if len(f.task.Metadata) > 0 && string(f.task.Metadata) != "null" {
		_ = json.Unmarshal(f.task.Metadata, &current)
	}
	if current == nil {
		current = make(map[string]any)
	}
	for k, v := range updates {
		current[k] = v
	}
	f.task.Metadata, _ = json.Marshal(current)
	return f.task, nil
}

func TestWorkCompletesThroughLifecycle(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Title:   "Test task",
			Source:  "test",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle: fakeLifecycle,
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	if len(fakeLifecycle.started) != 1 || fakeLifecycle.started[0] != "task-1" {
		t.Fatalf("expected started to have task-1, got %v", fakeLifecycle.started)
	}
	if len(fakeLifecycle.completed) != 1 || fakeLifecycle.completed[0] != "task-1" {
		t.Fatalf("expected completed to have task-1, got %v", fakeLifecycle.completed)
	}
	if len(fakeLifecycle.failed) != 0 {
		t.Fatalf("expected no failed tasks, got %v", fakeLifecycle.failed)
	}
}

func TestWorkFailsThroughLifecycle(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Title:   "Test task",
			Source:  "test",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle: fakeLifecycle,
		Agent: fakeAgentClient{
			err: errors.New("agent failed"),
		},
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to return error")
	}

	if len(fakeLifecycle.started) != 1 || fakeLifecycle.started[0] != "task-1" {
		t.Fatalf("expected started to have task-1, got %v", fakeLifecycle.started)
	}
	if len(fakeLifecycle.completed) != 0 {
		t.Fatalf("expected no completed tasks, got %v", fakeLifecycle.completed)
	}
	if len(fakeLifecycle.failed) != 1 || fakeLifecycle.failed[0] != "task-1" {
		t.Fatalf("expected failed to have task-1, got %v", fakeLifecycle.failed)
	}
}

type blockingModelClient struct{}

func (m blockingModelClient) Generate(ctx context.Context, input model.GenerateInput) (model.GenerateResult, error) {
	<-ctx.Done()
	return model.GenerateResult{}, ctx.Err()
}

func TestWorkMarksTimeoutFailure(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Title:   "Test task",
			Source:  "test",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle:  fakeLifecycle,
		Model:      blockingModelClient{},
		RunTimeout: 1 * time.Millisecond,
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to return timeout error")
	}
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("expected context.DeadlineExceeded error, got: %v", err)
	}

	if len(fakeLifecycle.started) != 1 || fakeLifecycle.started[0] != "task-1" {
		t.Fatalf("expected started to have task-1, got %v", fakeLifecycle.started)
	}
	if len(fakeLifecycle.completed) != 0 {
		t.Fatalf("expected no completed tasks, got %v", fakeLifecycle.completed)
	}
	if len(fakeLifecycle.failed) != 1 || fakeLifecycle.failed[0] != "task-1" {
		t.Fatalf("expected failed to have task-1, got %v", fakeLifecycle.failed)
	}
	if len(fakeLifecycle.failMessages) != 1 || !strings.Contains(fakeLifecycle.failMessages[0], "timeout") {
		t.Fatalf("expected fail message containing timeout, got %v", fakeLifecycle.failMessages)
	}
	if len(fakeLifecycle.failInputs) != 1 {
		t.Fatalf("expected 1 fail input, got %d", len(fakeLifecycle.failInputs))
	}
	if fakeLifecycle.failInputs[0].Type != workflow.FailureTypeTimeout {
		t.Errorf("expected failure type %q, got %q", workflow.FailureTypeTimeout, fakeLifecycle.failInputs[0].Type)
	}
}

func TestWorkRetriesAfterFailureState(t *testing.T) {
	sink := &spySink{}
	notifService := notification.NewService(nil, sink)

	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Status:  "pending",
			Title:   "Test task",
			Source:  "test",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle:     fakeLifecycle,
		Notifications: notifService,
		Agent: fakeAgentClient{
			err: errors.New("first attempt agent failure"),
		},
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected first Work to return error")
	}

	if fakeLifecycle.task.Status != "failed" {
		t.Fatalf("expected task status to be failed, got: %s", fakeLifecycle.task.Status)
	}

	if len(fakeLifecycle.failInputs) != 1 {
		t.Fatalf("expected 1 fail input, got %d", len(fakeLifecycle.failInputs))
	}
	if fakeLifecycle.failInputs[0].Type != workflow.FailureTypeExecution {
		t.Errorf("expected failure type %q, got %q", workflow.FailureTypeExecution, fakeLifecycle.failInputs[0].Type)
	}

	worker.Agent = nil

	err = worker.Work(context.Background(), job)
	if err != nil {
		t.Fatalf("expected second Work to succeed, got error: %v", err)
	}

	if fakeLifecycle.task.Status != "completed" {
		t.Fatalf("expected task status to be completed, got: %s", fakeLifecycle.task.Status)
	}

	if len(fakeLifecycle.started) != 2 {
		t.Fatalf("expected StartTask to be called 2 times, got: %d", len(fakeLifecycle.started))
	}
	if len(fakeLifecycle.failed) != 1 {
		t.Fatalf("expected FailTask to be called 1 time, got: %d", len(fakeLifecycle.failed))
	}
	if len(fakeLifecycle.completed) != 1 {
		t.Fatalf("expected CompleteTask to be called 1 time, got: %d", len(fakeLifecycle.completed))
	}

	// Verify the emitted notification event sequence and context fields (preventing stale reason leak)
	if len(sink.events) != 4 {
		t.Fatalf("expected 4 notification events fanned out, got %d: %+v", len(sink.events), sink.events)
	}

	// First attempt Running event
	ev0 := sink.events[0]
	if ev0.Type != notification.TaskEventRunning {
		t.Errorf("ev0: expected type Running, got %s", ev0.Type)
	}
	if ev0.Attempt != 1 {
		t.Errorf("ev0: expected attempt 1, got %d", ev0.Attempt)
	}
	if ev0.Reason != "" {
		t.Errorf("ev0: expected empty reason (no stale failure leak), got %q", ev0.Reason)
	}

	// First attempt Failed event
	ev1 := sink.events[1]
	if ev1.Type != notification.TaskEventFailed {
		t.Errorf("ev1: expected type Failed, got %s", ev1.Type)
	}
	if ev1.Attempt != 1 {
		t.Errorf("ev1: expected attempt 1, got %d", ev1.Attempt)
	}
	if ev1.Reason != "first attempt agent failure" {
		t.Errorf("ev1: expected reason 'first attempt agent failure', got %q", ev1.Reason)
	}

	// Second attempt Running event
	ev2 := sink.events[2]
	if ev2.Type != notification.TaskEventRunning {
		t.Errorf("ev2: expected type Running, got %s", ev2.Type)
	}
	if ev2.Attempt != 2 {
		t.Errorf("ev2: expected attempt 2, got %d", ev2.Attempt)
	}
	if ev2.Reason != "" {
		t.Errorf("ev2: expected empty reason (prevented stale failed reason leak), got %q", ev2.Reason)
	}

	// Second attempt Completed event
	ev3 := sink.events[3]
	if ev3.Type != notification.TaskEventCompleted {
		t.Errorf("ev3: expected type Completed, got %s", ev3.Type)
	}
	if ev3.Attempt != 2 {
		t.Errorf("ev3: expected attempt 2, got %d", ev3.Attempt)
	}
	if ev3.Reason != "" {
		t.Errorf("ev3: expected empty reason (prevented stale failed reason leak), got %q", ev3.Reason)
	}
}

func checkoutTaskMeta(slotPath, sourceBranch string) json.RawMessage {
	raw, _ := json.Marshal(map[string]any{
		"checkout": map[string]any{
			"slot_path":     slotPath,
			"source_branch": sourceBranch,
			"slot_index":    0,
			"slot_id":       1,
		},
	})
	return raw
}

func stringPtr(s string) *string {
	return &s
}

func withPlaneExternalRef(task storage.Task, workItemID string) storage.Task {
	task.ExternalProvider = stringPtr("plane")
	task.ExternalID = stringPtr(workItemID)
	return task
}

func TestRunTaskRoutesPlaneOriginToAuthoringBridge(t *testing.T) {
	var capturedInput model.GenerateInput

	fakeModel := fakeModelClient{
		generate: func(_ context.Context, input model.GenerateInput) (model.GenerateResult, error) {
			capturedInput = input
			return model.GenerateResult{
				ID:    "resp-1",
				Model: "m",
				Text:  "authoring done",
			}, nil
		},
	}

	worker := &TaskWorker{Model: fakeModel, IdentityWriter: &fakeIdentityWriter{}}

	task := withPlaneExternalRef(storage.Task{
		ID:       "task-plane-1",
		Title:    "NOMAD-42: Create milestone",
		Source:   "plane",
		Payload:  json.RawMessage(`{"prompt":"Author milestone for NOMAD-42","instructions":"Use roadmap skill"}`),
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}, "NOMAD-42")

	raw, msg, err := worker.runTask(context.Background(), task)
	if err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}
	if msg == "" {
		t.Fatal("expected non-empty message")
	}
	if raw == nil {
		t.Fatal("expected non-nil result")
	}

	if capturedInput.WorkspaceMetadata == nil {
		t.Fatal("expected WorkspaceMetadata to be set for Plane-origin authoring task")
	}
	if capturedInput.WorkspaceMetadata.Path != "/home/user/workspace/nomadcode/slots/000" {
		t.Errorf("unexpected workspace path: %q", capturedInput.WorkspaceMetadata.Path)
	}
	if capturedInput.WorkspaceMetadata.SourceBranch != "develop" {
		t.Errorf("unexpected source branch: %q", capturedInput.WorkspaceMetadata.SourceBranch)
	}
	if capturedInput.WorkspaceMetadata.Provider != "plane" {
		t.Errorf("unexpected provider: %q", capturedInput.WorkspaceMetadata.Provider)
	}
	if capturedInput.WorkspaceMetadata.WorkItemID != "NOMAD-42" {
		t.Errorf("unexpected work_item_id: %q", capturedInput.WorkspaceMetadata.WorkItemID)
	}
}

func TestRunTaskAuthoringMissingIdentityReturnsExplicitError(t *testing.T) {
	var modelCalled bool
	worker := &TaskWorker{Model: fakeModelClient{
		generate: func(context.Context, model.GenerateInput) (model.GenerateResult, error) {
			modelCalled = true
			return model.GenerateResult{}, nil
		},
	}}

	task := storage.Task{
		ID:       "task-missing-identity",
		Title:    "NOMAD-42: Create milestone",
		Source:   "plane",
		Payload:  json.RawMessage(`{"prompt":"Author milestone for NOMAD-42"}`),
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}

	_, _, err := worker.runTask(context.Background(), task)
	if !errors.Is(err, authoring.ErrProviderIdentityMissing) {
		t.Fatalf("expected ErrProviderIdentityMissing, got %v", err)
	}
	if modelCalled {
		t.Fatal("model must not be called when provider identity is missing")
	}
}

func TestRunTaskAuthoringPreemptsA2AClient(t *testing.T) {
	var modelCalled bool
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, input model.GenerateInput) (model.GenerateResult, error) {
			modelCalled = true
			if input.WorkspaceMetadata == nil {
				t.Error("expected WorkspaceMetadata for authoring task routed through model")
			}
			return model.GenerateResult{Text: "done", Model: "m"}, nil
		},
	}
	// A2A client is set — authoring task must still go to model bridge.
	worker := &TaskWorker{
		Model:          fakeModel,
		Agent:          fakeAgentClient{err: errors.New("a2a should not be called")},
		IdentityWriter: &fakeIdentityWriter{},
	}

	task := withPlaneExternalRef(storage.Task{
		ID:       "task-preempt",
		Source:   "plane",
		Payload:  json.RawMessage(`{}`),
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}, "NOMAD-42")

	raw, _, err := worker.runTask(context.Background(), task)
	if err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}
	if !modelCalled {
		t.Error("expected model client to be called for authoring task, not A2A")
	}
	var result map[string]any
	_ = json.Unmarshal(raw, &result)
	if result["mode"] != "authoring_run" {
		t.Errorf("expected mode=authoring_run, got %#v", result["mode"])
	}
}

func TestRunTaskGenericModelHasNoWorkspaceMetadata(t *testing.T) {
	var capturedInput model.GenerateInput

	fakeModel := fakeModelClient{
		generate: func(_ context.Context, input model.GenerateInput) (model.GenerateResult, error) {
			capturedInput = input
			return model.GenerateResult{Text: "done"}, nil
		},
	}

	worker := &TaskWorker{Model: fakeModel}

	task := storage.Task{
		ID:      "task-manual-1",
		Title:   "Generic task",
		Source:  "manual",
		Payload: json.RawMessage(`{"prompt":"do something"}`),
	}

	if _, _, err := worker.runTask(context.Background(), task); err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}
	if capturedInput.WorkspaceMetadata != nil {
		t.Errorf("expected no WorkspaceMetadata for generic task, got %#v", capturedInput.WorkspaceMetadata)
	}
}

type fakeModelClient struct {
	generate func(context.Context, model.GenerateInput) (model.GenerateResult, error)
}

func (f fakeModelClient) Generate(ctx context.Context, input model.GenerateInput) (model.GenerateResult, error) {
	return f.generate(ctx, input)
}

func TestRunTaskAuthoringSuccessRecordsState(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{Text: "milestone authored", Model: "m"}, nil
		},
	}
	worker := &TaskWorker{Model: fakeModel, IdentityWriter: &fakeIdentityWriter{}}

	task := withPlaneExternalRef(storage.Task{
		ID:       "task-authoring-ok",
		Title:    "Author milestone",
		Source:   "plane",
		Payload:  json.RawMessage(`{}`),
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}, "NOMAD-42")

	raw, _, err := worker.runTask(context.Background(), task)
	if err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}

	var result map[string]any
	if err := json.Unmarshal(raw, &result); err != nil {
		t.Fatalf("unmarshal result: %v", err)
	}
	if result["authoring_run_state"] != "in_progress" {
		t.Errorf("expected authoring_run_state=in_progress, got %#v", result["authoring_run_state"])
	}
	if result["authoring_run_updated_at"] == nil || result["authoring_run_updated_at"] == "" {
		t.Error("expected authoring_run_updated_at to be set")
	}
	if result["mode"] != "authoring_run" {
		t.Errorf("expected mode=authoring_run, got %#v", result["mode"])
	}
}

func TestWorkRecordsAuthoringInProgressBeforeExecution(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{Text: "done", Model: "m"}, nil
		},
	}

	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-ip",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		IdentityWriter: &fakeIdentityWriter{},
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-ip"}}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	// MergeTaskMetadata must have been called at least once with in_progress.
	found := false
	for _, m := range fakeLifecycle.mergedMetadata {
		if m[workflow.MetadataKeyAuthoringRunState] == "in_progress" {
			found = true
			if m[workflow.MetadataKeyAuthoringRunUpdatedAt] == nil {
				t.Error("expected authoring_run_updated_at to be set with in_progress")
			}
		}
	}
	if !found {
		t.Errorf("expected MergeTaskMetadata to be called with authoring_run_state=in_progress, got %v", fakeLifecycle.mergedMetadata)
	}
}

func TestWorkDoesNotRecordAuthoringInProgressForGenericTask(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:     "task-generic-ip",
			Source: "manual",
			Status: "pending",
		},
	}

	worker := &TaskWorker{Lifecycle: fakeLifecycle}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-generic-ip"}}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	for _, m := range fakeLifecycle.mergedMetadata {
		if m[workflow.MetadataKeyAuthoringRunState] != nil {
			t.Errorf("expected no authoring metadata for generic task, got %v", m)
		}
	}
}

func TestMarkFailedWithTaskAuthoringRecordsState(t *testing.T) {
	lc := &fakeTaskLifecycle{
		task: storage.Task{
			ID:       "task-authoring-fail",
			Title:    "Author milestone",
			Source:   "plane",
			Status:   "running",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		},
	}
	lc.task.Metadata = checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop")

	slotUpdater := &fakeWorkspaceSlotUpdater{}
	worker := &TaskWorker{Lifecycle: lc, SlotUpdater: slotUpdater}

	task := storage.Task{
		ID:       "task-authoring-fail",
		Source:   "plane",
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}

	worker.markFailedWithTask("task-authoring-fail", task, errors.New("IOP bridge failed"))

	if len(lc.failInputs) != 1 {
		t.Fatalf("expected 1 fail input, got %d", len(lc.failInputs))
	}
	extra := lc.failInputs[0].ExtraMetadata
	if extra == nil {
		t.Fatal("expected ExtraMetadata to be set for authoring task failure")
	}
	if extra[workflow.MetadataKeyAuthoringRunState] != "failed" {
		t.Errorf("expected authoring_run_state=failed, got %#v", extra[workflow.MetadataKeyAuthoringRunState])
	}
	if extra[workflow.MetadataKeyAuthoringRunUpdatedAt] == nil {
		t.Error("expected authoring_run_updated_at to be set")
	}
	if extra[workflow.MetadataKeyAuthoringFailureCategory] != "bridge_failed" {
		t.Errorf("expected authoring_failure_category=bridge_failed, got %#v", extra[workflow.MetadataKeyAuthoringFailureCategory])
	}
	if len(slotUpdater.calls) != 1 {
		t.Fatalf("expected 1 slot state update, got %d", len(slotUpdater.calls))
	}
	if slotUpdater.calls[0].ID != 1 || slotUpdater.calls[0].State != string(projectsync.SlotStateError) {
		t.Errorf("expected slot 1 -> error, got %#v", slotUpdater.calls[0])
	}
}

func TestMarkFailedWithTaskAuthoringMissingSlotIDSkipsSlotUpdate(t *testing.T) {
	metadata, _ := json.Marshal(map[string]any{
		"checkout": map[string]any{
			"slot_path":     "/home/user/workspace/nomadcode/slots/000",
			"source_branch": "develop",
		},
	})
	lc := &fakeTaskLifecycle{
		task: storage.Task{
			ID:       "task-authoring-missing-slot-id",
			Source:   "plane",
			Status:   "running",
			Metadata: metadata,
		},
	}
	slotUpdater := &fakeWorkspaceSlotUpdater{}
	worker := &TaskWorker{Lifecycle: lc, SlotUpdater: slotUpdater}

	task := storage.Task{
		ID:       "task-authoring-missing-slot-id",
		Source:   "plane",
		Metadata: metadata,
	}

	worker.markFailedWithTask("task-authoring-missing-slot-id", task, errors.New("IOP bridge failed"))

	if len(lc.failInputs) != 1 {
		t.Fatalf("expected task failure to be recorded, got %d fail inputs", len(lc.failInputs))
	}
	if len(slotUpdater.calls) != 0 {
		t.Fatalf("expected no slot update without slot_id, got %#v", slotUpdater.calls)
	}
}

func TestMarkFailedWithTaskGenericNoAuthoringState(t *testing.T) {
	lc := &fakeTaskLifecycle{
		task: storage.Task{
			ID:     "task-generic-fail",
			Source: "manual",
			Status: "running",
		},
	}
	worker := &TaskWorker{Lifecycle: lc}

	task := storage.Task{
		ID:     "task-generic-fail",
		Source: "manual",
	}

	worker.markFailedWithTask("task-generic-fail", task, errors.New("generic error"))

	if len(lc.failInputs) != 1 {
		t.Fatalf("expected 1 fail input, got %d", len(lc.failInputs))
	}
	if lc.failInputs[0].ExtraMetadata != nil {
		t.Errorf("expected no ExtraMetadata for generic task, got %#v", lc.failInputs[0].ExtraMetadata)
	}
}

type logRecord struct {
	msg  string
	args map[string]any
}

type spyHandler struct {
	records []logRecord
}

func (h *spyHandler) Enabled(ctx context.Context, level slog.Level) bool {
	return true
}

func (h *spyHandler) Handle(ctx context.Context, r slog.Record) error {
	args := make(map[string]any)
	r.Attrs(func(a slog.Attr) bool {
		args[a.Key] = a.Value.Any()
		return true
	})
	h.records = append(h.records, logRecord{
		msg:  r.Message,
		args: args,
	})
	return nil
}

func (h *spyHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return h
}

func (h *spyHandler) WithGroup(name string) slog.Handler {
	return h
}

type spySink struct {
	events []notification.TaskEvent
}

func (s *spySink) HandleTaskEvent(ctx context.Context, event notification.TaskEvent) error {
	s.events = append(s.events, event)
	return nil
}

func TestWorkEmitsRunningAndCompletedEvents(t *testing.T) {
	spy := &spyHandler{}
	logger := slog.New(spy)
	sink := &spySink{}
	notifService := notification.NewService(logger, sink)

	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-1",
			Title:   "Success Task",
			Source:  "test",
			Status:  "pending",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle:     fakeLifecycle,
		Notifications: notifService,
		Logger:        logger,
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-1"},
	}

	err := worker.Work(context.Background(), job)
	if err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	// Verify events from sink
	if len(sink.events) != 2 {
		t.Fatalf("expected 2 notification events from sink, got %d", len(sink.events))
	}

	// Running Event
	ev0 := sink.events[0]
	if ev0.Type != notification.TaskEventRunning {
		t.Errorf("expected first event to be %s, got %s", notification.TaskEventRunning, ev0.Type)
	}
	if ev0.Attempt != 1 {
		t.Errorf("expected running attempt 1, got %d", ev0.Attempt)
	}
	if ev0.OccurredAt.IsZero() {
		t.Error("expected running occurred_at to be non-zero")
	}

	// Completed Event
	ev1 := sink.events[1]
	if ev1.Type != notification.TaskEventCompleted {
		t.Errorf("expected second event to be %s, got %s", notification.TaskEventCompleted, ev1.Type)
	}
	if ev1.Attempt != 1 {
		t.Errorf("expected completed attempt 1, got %d", ev1.Attempt)
	}
	if ev1.OccurredAt.IsZero() {
		t.Error("expected completed occurred_at to be non-zero")
	}
}

func TestWorkEmitsFailedEvent(t *testing.T) {
	spy := &spyHandler{}
	logger := slog.New(spy)
	sink := &spySink{}
	notifService := notification.NewService(logger, sink)

	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:      "task-2",
			Title:   "Failure Task",
			Source:  "test",
			Status:  "pending",
			Payload: json.RawMessage(`{}`),
		},
	}

	worker := &TaskWorker{
		Lifecycle:     fakeLifecycle,
		Notifications: notifService,
		Agent: fakeAgentClient{
			err: errors.New("agent failed"),
		},
		Logger: logger,
	}

	job := &river.Job[TaskJobArgs]{
		Args: TaskJobArgs{TaskID: "task-2"},
	}

	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to return error")
	}

	// Verify events from sink
	if len(sink.events) != 2 {
		t.Fatalf("expected 2 notification events from sink, got %d", len(sink.events))
	}

	// Running Event
	ev0 := sink.events[0]
	if ev0.Type != notification.TaskEventRunning {
		t.Errorf("expected first event to be %s, got %s", notification.TaskEventRunning, ev0.Type)
	}

	// Failed Event
	ev1 := sink.events[1]
	if ev1.Type != notification.TaskEventFailed {
		t.Errorf("expected second event to be %s, got %s", notification.TaskEventFailed, ev1.Type)
	}
	if ev1.Attempt != 1 {
		t.Errorf("expected failed attempt 1, got %d", ev1.Attempt)
	}
	if ev1.Reason != "agent failed" {
		t.Errorf("expected failed reason %q, got %q", "agent failed", ev1.Reason)
	}
	if ev1.OccurredAt.IsZero() {
		t.Error("expected failed occurred_at to be non-zero")
	}
}

func TestWorkAuthoringSkipsGenericRunTimeout(t *testing.T) {
	// done is closed after RunTimeout would have fired. If the scheduler applied
	// RunTimeout to authoring tasks the context would expire and this test would
	// fail with DeadlineExceeded.
	done := make(chan struct{})
	fakeModel := fakeModelClient{
		generate: func(ctx context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			select {
			case <-done:
				return model.GenerateResult{Text: "authoring done", Model: "m"}, nil
			case <-ctx.Done():
				return model.GenerateResult{}, ctx.Err()
			}
		},
	}

	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-timeout",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		RunTimeout:     1 * time.Millisecond,
		IdentityWriter: &fakeIdentityWriter{},
	}

	go func() {
		time.Sleep(20 * time.Millisecond)
		close(done)
	}()

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-timeout"}}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("authoring task with short RunTimeout returned error: %v", err)
	}
	if len(fakeLifecycle.failed) != 0 {
		t.Fatalf("expected authoring task not to be failed by generic RunTimeout, got failed: %v", fakeLifecycle.failed)
	}
}

func TestWorkKeepsAuthoringTaskRunningUntilDevelopMatch(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{Text: "milestone authored", Model: "m"}, nil
		},
	}

	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-wait",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		IdentityWriter: &fakeIdentityWriter{},
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-wait"}}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	if len(fakeLifecycle.completed) != 0 {
		t.Fatalf("expected CompletedTask to NOT be called, but was called: %v", fakeLifecycle.completed)
	}

	foundWaitType := false
	foundRunState := false
	for _, m := range fakeLifecycle.mergedMetadata {
		if m[workflow.MetadataKeyWaitType] == "develop_match" {
			foundWaitType = true
		}
		if m[workflow.MetadataKeyAuthoringRunState] == "in_progress" {
			foundRunState = true
		}
	}
	if !foundWaitType {
		t.Errorf("expected MergeTaskMetadata to be called with wait_type=develop_match, got %v", fakeLifecycle.mergedMetadata)
	}
	if !foundRunState {
		t.Errorf("expected MergeTaskMetadata to be called with authoring_run_state=in_progress, got %v", fakeLifecycle.mergedMetadata)
	}
}

func TestWorkAuthoringTwoCallsAndPushCommand(t *testing.T) {
	var calls []model.GenerateInput
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, input model.GenerateInput) (model.GenerateResult, error) {
			calls = append(calls, input)
			return model.GenerateResult{Text: "done", ID: fmt.Sprintf("resp-%d", len(calls)), Model: "m"}, nil
		},
	}

	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-two-calls",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		IdentityWriter: &fakeIdentityWriter{},
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-two-calls"}}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	if len(calls) != 2 {
		t.Fatalf("expected exactly 2 model calls (authoring + push-only), got %d", len(calls))
	}
	if calls[1].Input != "변경된 내용에 대해 develop 브런치에 푸시해" {
		t.Errorf("expected push-only command text, got %q", calls[1].Input)
	}
	if calls[1].WorkspaceMetadata == nil {
		t.Fatal("expected WorkspaceMetadata on push-only call")
	}
	if calls[1].WorkspaceMetadata.Path != "/home/user/workspace/nomadcode/slots/000" {
		t.Errorf("unexpected push-only workspace path: %q", calls[1].WorkspaceMetadata.Path)
	}

	// Completed must NOT be called; task stays running for develop match.
	if len(fakeLifecycle.completed) != 0 {
		t.Fatalf("expected CompleteTask to NOT be called, but was called: %v", fakeLifecycle.completed)
	}

	// wait_type=develop_match must be set only after push-only succeeds.
	foundWaitType := false
	for _, m := range fakeLifecycle.mergedMetadata {
		if m[workflow.MetadataKeyWaitType] == "develop_match" {
			foundWaitType = true
		}
	}
	if !foundWaitType {
		t.Errorf("expected MergeTaskMetadata to be called with wait_type=develop_match, got %v", fakeLifecycle.mergedMetadata)
	}
}

func TestWorkAuthoringPushFailureMarksPushFailed(t *testing.T) {
	callCount := 0
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			callCount++
			if callCount == 2 {
				return model.GenerateResult{}, errors.New("git push rejected")
			}
			return model.GenerateResult{Text: "authored", ID: "resp-1", Model: "m"}, nil
		},
	}

	fakeSlotUpdater := &fakeWorkspaceSlotUpdater{}
	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-push-fail",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		SlotUpdater:    fakeSlotUpdater,
		IdentityWriter: &fakeIdentityWriter{},
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-push-fail"}}
	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to return error on push failure")
	}

	if len(fakeLifecycle.failed) == 0 {
		t.Fatal("expected FailTask to be called on push failure")
	}
	if len(fakeLifecycle.failInputs) == 0 {
		t.Fatal("expected failInputs to be recorded")
	}

	fi := fakeLifecycle.failInputs[0]
	if fi.ExtraMetadata == nil {
		t.Fatal("expected ExtraMetadata on failure input")
	}
	if fi.ExtraMetadata["authoring_failure_category"] != "push_failed" {
		t.Errorf("expected authoring_failure_category=push_failed, got %#v", fi.ExtraMetadata["authoring_failure_category"])
	}
	if fi.ExtraMetadata["authoring_run_state"] != "failed" {
		t.Errorf("expected authoring_run_state=failed, got %#v", fi.ExtraMetadata["authoring_run_state"])
	}

	// Slot must be marked dirty.
	if len(fakeSlotUpdater.calls) == 0 {
		t.Fatal("expected slot state update on push failure")
	}
	if fakeSlotUpdater.calls[0].State != string(projectsync.SlotStateDirty) {
		t.Errorf("expected slot state dirty, got %q", fakeSlotUpdater.calls[0].State)
	}
}

func TestRunTaskAuthoringRecordsPushResponse(t *testing.T) {
	callCount := 0
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			callCount++
			return model.GenerateResult{ID: fmt.Sprintf("id-%d", callCount), Text: "done", Model: "m"}, nil
		},
	}
	worker := &TaskWorker{Model: fakeModel, IdentityWriter: &fakeIdentityWriter{}}

	task := withPlaneExternalRef(storage.Task{
		ID:       "task-push-resp",
		Source:   "plane",
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}, "NOMAD-42")

	raw, _, err := worker.runTask(context.Background(), task)
	if err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}

	var result map[string]any
	if err := json.Unmarshal(raw, &result); err != nil {
		t.Fatalf("unmarshal result: %v", err)
	}
	if result["push_response_id"] == "" || result["push_response_id"] == nil {
		t.Errorf("expected push_response_id in result, got %#v", result["push_response_id"])
	}
	if result["push_request_state"] != "succeeded" {
		t.Errorf("expected push_request_state=succeeded, got %#v", result["push_request_state"])
	}
}

func TestRunTaskAuthoringProgressCallback(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: storage.Task{
			ID:       "task-auth-progress",
			Source:   "plane",
			Status:   "running",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		},
	}

	progressTime := time.Date(2026, 6, 20, 12, 0, 0, 0, time.UTC)
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, input model.GenerateInput) (model.GenerateResult, error) {
			if input.OnProgress != nil {
				input.OnProgress(model.GenerateProgress{
					Mode:          "streaming",
					Reason:        "receiving stream chunks",
					LastEventTime: progressTime,
				})
			}
			return model.GenerateResult{
				ID:    "resp-progress",
				Model: "m",
				Text:  "authoring progress completed",
			}, nil
		},
	}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		IdentityWriter: &fakeIdentityWriter{},
	}

	task := withPlaneExternalRef(storage.Task{
		ID:       "task-auth-progress",
		Source:   "plane",
		Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
	}, "NOMAD-42")

	raw, msg, err := worker.runTask(context.Background(), task)
	if err != nil {
		t.Fatalf("runTask returned error: %v", err)
	}
	if msg == "" {
		t.Fatal("expected non-empty message")
	}

	foundProgress := false
	for _, m := range fakeLifecycle.mergedMetadata {
		if m[workflow.MetadataKeyAuthoringProgressMode] == "streaming" {
			foundProgress = true
			if m[workflow.MetadataKeyAuthoringProgressReason] != "receiving stream chunks" {
				t.Errorf("unexpected progress reason: %v", m[workflow.MetadataKeyAuthoringProgressReason])
			}
			if m[workflow.MetadataKeyAuthoringRunUpdatedAt] != progressTime.Format(time.RFC3339) {
				t.Errorf("unexpected updated at: %v", m[workflow.MetadataKeyAuthoringRunUpdatedAt])
			}
		}
	}
	if !foundProgress {
		t.Errorf("expected MergeTaskMetadata to be called with streaming progress, got %v", fakeLifecycle.mergedMetadata)
	}

	var result map[string]any
	if err := json.Unmarshal(raw, &result); err != nil {
		t.Fatalf("unmarshal result: %v", err)
	}
	if result[workflow.MetadataKeyAuthoringProgressMode] != "streaming" {
		t.Errorf("expected progress mode streaming in raw result, got %v", result[workflow.MetadataKeyAuthoringProgressMode])
	}
	if result[workflow.MetadataKeyAuthoringProgressReason] != "receiving stream chunks" {
		t.Errorf("expected progress reason in raw result, got %v", result[workflow.MetadataKeyAuthoringProgressReason])
	}
}

type fakeIdentityWriter struct {
	calls []authoring.EnsureProviderIdentityInput
	err   error
}

func (w *fakeIdentityWriter) EnsureProviderIdentity(ctx context.Context, in authoring.EnsureProviderIdentityInput) (authoring.EnsureProviderIdentityResult, error) {
	w.calls = append(w.calls, in)
	if w.err != nil {
		return authoring.EnsureProviderIdentityResult{}, w.err
	}
	return authoring.EnsureProviderIdentityResult{InjectedPaths: []string{"test-milestone.md"}}, nil
}

func TestWorkAuthoringWritesIdentityBeforePush(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{ID: "resp-1", Text: "done", Model: "m"}, nil
		},
	}
	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-ident-write",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}
	fakeSlotUpdater := &fakeWorkspaceSlotUpdater{}

	fakeWriter := &fakeIdentityWriter{}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		SlotUpdater:    fakeSlotUpdater,
		IdentityWriter: fakeWriter,
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-ident-write"}}
	err := worker.Work(context.Background(), job)
	if err != nil {
		t.Fatalf("unexpected Work error: %v", err)
	}

	if len(fakeWriter.calls) != 1 {
		t.Fatalf("expected 1 EnsureProviderIdentity call, got %d", len(fakeWriter.calls))
	}
	call := fakeWriter.calls[0]
	if call.WorkspacePath != "/home/user/workspace/nomadcode/slots/000" {
		t.Errorf("expected WorkspacePath=/home/user/workspace/nomadcode/slots/000, got %q", call.WorkspacePath)
	}
	if call.Identity.WorkItemID != "NOMAD-42" {
		t.Errorf("expected WorkItemID=NOMAD-42, got %q", call.Identity.WorkItemID)
	}
}

func TestWorkAuthoringIdentityWriteFailureMarksIdentityMissing(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{ID: "resp-1", Text: "done", Model: "m"}, nil
		},
	}
	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-ident-fail",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}
	fakeSlotUpdater := &fakeWorkspaceSlotUpdater{}

	// Case 1: Invalid identity -> IdentityMissing
	fakeWriter := &fakeIdentityWriter{err: roadmapsync.ErrInvalidIdentity}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		SlotUpdater:    fakeSlotUpdater,
		IdentityWriter: fakeWriter,
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-ident-fail"}}
	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to fail when identity writer fails")
	}

	if len(fakeLifecycle.failInputs) == 0 {
		t.Fatal("expected failInputs to be recorded")
	}
	fi := fakeLifecycle.failInputs[0]
	if fi.ExtraMetadata["authoring_failure_category"] != "identity_missing" {
		t.Errorf("expected authoring_failure_category=identity_missing, got %q", fi.ExtraMetadata["authoring_failure_category"])
	}
	if fi.ExtraMetadata["authoring_run_state"] != "failed" {
		t.Errorf("expected authoring_run_state=failed, got %q", fi.ExtraMetadata["authoring_run_state"])
	}
	if fakeSlotUpdater.calls[0].State != string(projectsync.SlotStateError) {
		t.Errorf("expected slot state error, got %q", fakeSlotUpdater.calls[0].State)
	}

	// Case 2: Other error -> IdentityWriteFailed
	fakeLifecycle2 := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-ident-fail-2",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}
	fakeSlotUpdater2 := &fakeWorkspaceSlotUpdater{}
	fakeWriter2 := &fakeIdentityWriter{err: errors.New("write permission denied")}
	worker2 := &TaskWorker{
		Lifecycle:      fakeLifecycle2,
		Model:          fakeModel,
		SlotUpdater:    fakeSlotUpdater2,
		IdentityWriter: fakeWriter2,
	}
	job2 := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-ident-fail-2"}}
	err2 := worker2.Work(context.Background(), job2)
	if err2 == nil {
		t.Fatal("expected Work to fail when identity writer fails with generic error")
	}
	if len(fakeLifecycle2.failInputs) == 0 {
		t.Fatal("expected failInputs to be recorded")
	}
	fi2 := fakeLifecycle2.failInputs[0]
	if fi2.ExtraMetadata["authoring_failure_category"] != "identity_write_failed" {
		t.Errorf("expected authoring_failure_category=identity_write_failed, got %q", fi2.ExtraMetadata["authoring_failure_category"])
	}
}

func TestSchedulerRegistersTaskWorkerWithIdentityWriter(t *testing.T) {
	fakeWriter := &fakeIdentityWriter{}
	client, err := New(lazyPool(t), nil, nil, nil, nil, fakeWriter, nil, nil, nil, 0, nil)
	if err != nil {
		t.Fatalf("New with identity writer failed: %v", err)
	}
	if client == nil {
		t.Fatal("expected non-nil client")
	}
}

func TestRunTaskAuthoringFailsWhenIdentityWriterMissing(t *testing.T) {
	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			return model.GenerateResult{ID: "resp-1", Text: "done", Model: "m"}, nil
		},
	}
	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-auth-missing-writer",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}
	fakeSlotUpdater := &fakeWorkspaceSlotUpdater{}

	// IdentityWriter is explicitly nil
	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		SlotUpdater:    fakeSlotUpdater,
		IdentityWriter: nil,
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-auth-missing-writer"}}
	err := worker.Work(context.Background(), job)
	if err == nil {
		t.Fatal("expected Work to fail when identity writer is missing")
	}

	if len(fakeLifecycle.failInputs) == 0 {
		t.Fatal("expected failInputs to be recorded")
	}
	fi := fakeLifecycle.failInputs[0]
	if fi.ExtraMetadata["authoring_failure_category"] != "identity_write_failed" {
		t.Errorf("expected authoring_failure_category=identity_write_failed, got %q", fi.ExtraMetadata["authoring_failure_category"])
	}
}

type fakeSyncRunner struct {
	result roadmapsyncpipeline.SyncCreationResult
	err    error
}

func (f fakeSyncRunner) SyncCreation(_ context.Context, _ roadmapsyncpipeline.SyncCreationInput) (roadmapsyncpipeline.SyncCreationResult, error) {
	return f.result, f.err
}

type fakeCreationTaskFinalizer struct {
	called bool
	err    error
}

func (f *fakeCreationTaskFinalizer) CompleteTaskByExternalRef(_ context.Context, _, _ string, _ json.RawMessage) (storage.Task, bool, error) {
	f.called = true
	return storage.Task{}, false, f.err
}

func TestRoadmapCreationSyncWorkerSkipsFinalizerWhenRefIDEmpty(t *testing.T) {
	// When SyncCreation returns SyncActionProjected but args.Ref.ID is empty
	// (missing-create case), TaskFinalizer must NOT be called to avoid
	// pgx.ErrNoRows causing an unintended River retry loop.
	finalizer := &fakeCreationTaskFinalizer{}
	worker := &RoadmapCreationSyncWorker{
		Sync:          fakeSyncRunner{result: roadmapsyncpipeline.SyncCreationResult{Action: roadmapsyncpipeline.SyncActionProjected}},
		TaskFinalizer: finalizer,
	}
	job := &river.Job[RoadmapCreationSyncJobArgs]{
		Args: RoadmapCreationSyncJobArgs{
			// Ref.ID is zero-value (empty) — missing-create case
		},
	}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}
	if finalizer.called {
		t.Error("TaskFinalizer.CompleteTaskByExternalRef must not be called when Ref.ID is empty (missing-create case)")
	}
}

func TestRoadmapCreationSyncWorkerCallsFinalizerWhenRefIDSet(t *testing.T) {
	// When SyncCreation returns SyncActionProjected and args.Ref.ID is non-empty
	// (reconcile case), TaskFinalizer IS called so workflow tasks can be completed.
	finalizer := &fakeCreationTaskFinalizer{}
	worker := &RoadmapCreationSyncWorker{
		Sync:          fakeSyncRunner{result: roadmapsyncpipeline.SyncCreationResult{Action: roadmapsyncpipeline.SyncActionProjected}},
		TaskFinalizer: finalizer,
	}
	job := &river.Job[RoadmapCreationSyncJobArgs]{
		Args: RoadmapCreationSyncJobArgs{
			Ref: workitem.Ref{
				Provider: "plane",
				ID:       "6945a5e7-b2ef-4d48-a037-00fa2b6ab99b",
			},
		},
	}
	if err := worker.Work(context.Background(), job); err != nil {
		t.Fatalf("Work returned error: %v", err)
	}
	if !finalizer.called {
		t.Error("TaskFinalizer.CompleteTaskByExternalRef must be called when Ref.ID is non-empty (reconcile case)")
	}
}

func TestTaskWorkerLateReturnMetadataMergeRace(t *testing.T) {
	fakeLifecycle := &fakeTaskLifecycle{
		task: withPlaneExternalRef(storage.Task{
			ID:       "task-race-1",
			Source:   "plane",
			Status:   "pending",
			Metadata: checkoutTaskMeta("/home/user/workspace/nomadcode/slots/000", "develop"),
		}, "NOMAD-42"),
	}

	fakeModel := fakeModelClient{
		generate: func(_ context.Context, _ model.GenerateInput) (model.GenerateResult, error) {
			// Simulate that external finalizer completed the task while TaskWorker is running.
			fakeLifecycle.task.Status = "completed"
			fakeLifecycle.task.Metadata = json.RawMessage(`{"authoring_run_state":"succeeded"}`)
			return model.GenerateResult{
				ID:    "resp-1",
				Text:  "milestone authored",
				Model: "m",
			}, nil
		},
	}

	fakeWriter := &fakeIdentityWriter{}

	worker := &TaskWorker{
		Lifecycle:      fakeLifecycle,
		Model:          fakeModel,
		IdentityWriter: fakeWriter,
	}

	job := &river.Job[TaskJobArgs]{Args: TaskJobArgs{TaskID: "task-race-1"}}
	err := worker.Work(context.Background(), job)
	if err != nil {
		t.Fatalf("Work returned error: %v", err)
	}

	// Verify that the terminal status remains completed and the metadata is not overwritten by in_progress/develop_match.
	if fakeLifecycle.task.Status != "completed" {
		t.Errorf("expected task status to remain completed, got %v", fakeLifecycle.task.Status)
	}

	var meta map[string]any
	if err := json.Unmarshal(fakeLifecycle.task.Metadata, &meta); err != nil {
		t.Fatalf("unmarshal metadata: %v", err)
	}

	if meta["authoring_run_state"] != "succeeded" {
		t.Errorf("expected authoring_run_state to remain succeeded, got %v", meta["authoring_run_state"])
	}
	if val, ok := meta["wait_type"]; ok && val != nil {
		t.Errorf("expected wait_type to remain unset/nil, got %v", val)
	}
}

func TestTaskWorkerTimeoutHook(t *testing.T) {
	worker := &TaskWorker{}
	timeout := worker.Timeout(nil)
	if timeout != 15*time.Minute {
		t.Errorf("expected Timeout to return 15m, got %v", timeout)
	}
}