package runmanage

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
	workflowrunstore "github.com/tidbcloud/AgentTeam/internal/store/workflow/run"
)

type authorizedRepositoryStub struct {
	repositoryStub
	allow          bool
	seen           workflowrunstore.AgentPrincipal
	authorizeCalls int
}

func (stub *authorizedRepositoryStub) AuthorizeAgentRun(_ context.Context, principal workflowrunstore.AgentPrincipal, _, _ string) error {
	stub.authorizeCalls++
	stub.seen = principal
	if !stub.allow {
		return storeerr.NotFound("workflow Run not found")
	}
	return nil
}

func TestAgentRunReadsRequireDurablePrincipalBoundary(t *testing.T) {
	now := time.Now().UTC()
	run := workflowrunstore.Run{
		ID: "wri_agent_read", WorkspaceID: "wsp_agent_read", InvocationID: "wiv_agent_read",
		PlanID: "pln_agent_read", PlanVersionID: "plv_agent_read", Status: workflowrunstore.RunStatusWaiting,
		InputJSON: []byte(`{}`), StateJSON: []byte(`{}`), QueuedAt: now, Revision: 1,
	}
	principal := workflowrunstore.AgentPrincipal{WorkspaceID: run.WorkspaceID, AgentID: "agt_authorized", MemberID: "mem_authorized"}

	missingBoundary := New(repositoryStub{loadRun: func(context.Context, string, string) (workflowrunstore.Run, error) {
		t.Fatal("workspace-only LoadRun must not be reached")
		return workflowrunstore.Run{}, nil
	}})
	_, err := missingBoundary.GetRunForAgent(context.Background(), run.WorkspaceID, run.ID, principal)
	require.Error(t, err)
	require.True(t, errors.Is(err, storeerr.ErrNotFound))

	repository := &authorizedRepositoryStub{
		repositoryStub: repositoryStub{loadRun: func(context.Context, string, string) (workflowrunstore.Run, error) { return run, nil }},
		allow:          true,
	}
	service := New(repository)
	loaded, err := service.GetRunForAgent(context.Background(), run.WorkspaceID, run.ID, principal)
	require.NoError(t, err)
	require.Equal(t, run.ID, loaded.ID)
	require.Equal(t, principal, repository.seen)
	require.Equal(t, 1, repository.authorizeCalls)

	repository.allow = false
	_, err = service.ReadArtifactPageForAgent(context.Background(), ArtifactReadInput{
		WorkspaceID: run.WorkspaceID, RunID: run.ID,
	}, principal)
	require.Error(t, err)
	require.True(t, errors.Is(err, storeerr.ErrNotFound))
	require.Equal(t, 2, repository.authorizeCalls, "each artifact page must reauthorize the authenticated Agent")
}