package taskloop

import (
	"context"
	"encoding/json"
	"os"
	"testing"

	"iop/packages/go/agenttask"
)

func TestRecoveryClassifiesSubmittedLiveExitedStaleAndAmbiguous(t *testing.T) {
	fixture := newRuntimeFixture(t)
	recovery := NewRecovery(fixture.snapshot, staticArtifactRoot(fixture.projectA))
	project := agenttask.ProjectRecord{
		ProjectID: "project-a", WorkspaceID: WorkspaceIdentity(fixture.projectA),
	}
	baseWork := agenttask.WorkRecord{
		Unit: agenttask.WorkUnit{
			ID: "1",
			Metadata: map[string]string{
				"review_path": "agent-task/m-m1/1_first/CODE_REVIEW-test.md",
			},
		},
		AttemptID: "attempt-1",
	}
	processIdentity := func(opaque string) agenttask.LocatorRecord {
		return agenttask.LocatorRecord{
			Kind: agenttask.LocatorProcess, Opaque: opaque, Revision: "r1",
			ProjectID: project.ProjectID, WorkspaceID: project.WorkspaceID,
			WorkUnitID: "1", AttemptID: baseWork.AttemptID,
		}
	}
	liveOpaque, err := processLocator(os.Getpid())
	if err != nil {
		t.Fatalf("processLocator(current): %v", err)
	}
	staleBytes, err := json.Marshal(processReference{PID: os.Getpid(), StartToken: "stale"})
	if err != nil {
		t.Fatal(err)
	}
	missingBytes, err := json.Marshal(processReference{PID: 99999999, StartToken: "missing"})
	if err != nil {
		t.Fatal(err)
	}
	submission := agenttask.Submission{
		ProjectID: project.ProjectID, WorkUnitID: "1",
		AttemptID: baseWork.AttemptID, ArtifactID: artifactIdentity(baseWork), Ready: true,
	}

	tests := []struct {
		name string
		work agenttask.WorkRecord
		loc  map[agenttask.LocatorKind]agenttask.LocatorRecord
		want agenttask.RecoveryExecutionState
	}{
		{
			name: "persisted submission",
			work: func() agenttask.WorkRecord {
				work := baseWork
				work.Submission = &submission
				return work
			}(),
			want: agenttask.RecoveryExecutionSubmitted,
		},
		{
			name: "live exact process",
			work: baseWork,
			loc: map[agenttask.LocatorKind]agenttask.LocatorRecord{
				agenttask.LocatorProcess: processIdentity(liveOpaque),
			},
			want: agenttask.RecoveryExecutionLive,
		},
		{
			name: "exited process with completed artifact",
			work: baseWork,
			loc: map[agenttask.LocatorKind]agenttask.LocatorRecord{
				agenttask.LocatorProcess: processIdentity(string(missingBytes)),
				agenttask.LocatorSession: {
					Kind: agenttask.LocatorSession, Opaque: "session", Revision: "session-r1",
					ProjectID: project.ProjectID, WorkspaceID: project.WorkspaceID,
					WorkUnitID: "1", AttemptID: baseWork.AttemptID,
				},
				agenttask.LocatorOverlay: {
					Kind: agenttask.LocatorOverlay, Opaque: "overlay", Revision: "overlay-r1",
					ProjectID: project.ProjectID, WorkspaceID: project.WorkspaceID,
					WorkUnitID: "1", AttemptID: baseWork.AttemptID,
				},
			},
			want: agenttask.RecoveryExecutionSubmitted,
		},
		{
			name: "stale reused process",
			work: baseWork,
			loc: map[agenttask.LocatorKind]agenttask.LocatorRecord{
				agenttask.LocatorProcess: processIdentity(string(staleBytes)),
			},
			want: agenttask.RecoveryExecutionAmbiguous,
		},
		{
			name: "malformed locator",
			work: baseWork,
			loc: map[agenttask.LocatorKind]agenttask.LocatorRecord{
				agenttask.LocatorProcess: processIdentity("{"),
			},
			want: agenttask.RecoveryExecutionAmbiguous,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			observation, err := recovery.Inspect(context.Background(), agenttask.RecoveryRequest{
				Project: project, Work: test.work, Locators: test.loc,
			})
			if err != nil {
				t.Fatalf("Inspect: %v", err)
			}
			if observation.Execution != test.want {
				t.Fatalf("execution = %s, want %s", observation.Execution, test.want)
			}
			if test.name == "exited process with completed artifact" {
				if observation.Submission == nil ||
					len(observation.Submission.Locators) != 2 {
					t.Fatalf("recovered submission locators = %#v", observation.Submission)
				}
				for _, locator := range observation.Submission.Locators {
					if locator.Kind != agenttask.LocatorProcess &&
						locator.Kind != agenttask.LocatorSession {
						t.Fatalf("recovered submission retained host locator %q", locator.Kind)
					}
				}
			}
		})
	}
}

func TestRecoveryReportsExitedForPlaceholderAndExactCompletion(t *testing.T) {
	fixture := newRuntimeFixture(t)
	recovery := NewRecovery(fixture.snapshot, staticArtifactRoot(fixture.projectA))
	project := agenttask.ProjectRecord{
		ProjectID: "project-a", WorkspaceID: WorkspaceIdentity(fixture.projectA),
	}
	work := agenttask.WorkRecord{
		Unit: agenttask.WorkUnit{
			ID: "1",
			Metadata: map[string]string{
				"review_path": "agent-task/m-m1/1_first/CODE_REVIEW-test.md",
			},
		},
		AttemptID: "attempt-1",
	}
	reviewPath := fixture.projectA + "/agent-task/m-m1/1_first/CODE_REVIEW-test.md"
	writeTaskFile(t, reviewPath, "_Paste actual stdout/stderr._\n")
	missing, _ := json.Marshal(processReference{PID: 99999999, StartToken: "missing"})
	process := agenttask.LocatorRecord{
		Kind: agenttask.LocatorProcess, Opaque: string(missing), Revision: "process-r1",
		ProjectID: project.ProjectID, WorkspaceID: project.WorkspaceID,
		WorkUnitID: "1", AttemptID: work.AttemptID,
	}
	completion := agenttask.LocatorRecord{
		Kind: agenttask.LocatorCompletion, Opaque: "integration:done", Revision: "completion-r1",
		ProjectID: project.ProjectID, WorkspaceID: project.WorkspaceID,
		WorkUnitID: "1", AttemptID: work.AttemptID,
	}
	work.Integration = &agenttask.IntegrationResult{
		Outcome:           agenttask.IntegrationOutcomeIntegrated,
		CompletionLocator: &completion,
	}
	observation, err := recovery.Inspect(context.Background(), agenttask.RecoveryRequest{
		Project: project,
		Work:    work,
		Locators: map[agenttask.LocatorKind]agenttask.LocatorRecord{
			agenttask.LocatorProcess:    process,
			agenttask.LocatorCompletion: completion,
		},
	})
	if err != nil {
		t.Fatalf("Inspect: %v", err)
	}
	if observation.Execution != agenttask.RecoveryExecutionExited ||
		observation.Completion != agenttask.RecoveryCompletionComplete {
		t.Fatalf("observation = %#v", observation)
	}
}