package agent

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

	"github.com/stretchr/testify/require"
	workflowrun "github.com/tidbcloud/AgentTeam/internal/api/workflow/run"
	workflowrunstore "github.com/tidbcloud/AgentTeam/internal/store/workflow/run"
	workflowpayload "github.com/tidbcloud/AgentTeam/internal/workflow/payload"
)

type largePayloadRunService struct {
	publicRouteRunService
	run workflowrunstore.Run
}

func (service largePayloadRunService) GetRun(_ context.Context, _, _ string) (workflowrunstore.Run, error) {
	return service.run, nil
}

func (service largePayloadRunService) GetRunForAgent(ctx context.Context, workspaceID, runID string, _ workflowrun.AgentPrincipal) (workflowrunstore.Run, error) {
	return service.GetRun(ctx, workspaceID, runID)
}

type mutableLargePayloadRunService struct {
	publicRouteRunService
	run workflowrunstore.Run
}

func (service *mutableLargePayloadRunService) GetRun(_ context.Context, _, _ string) (workflowrunstore.Run, error) {
	return service.run, nil
}

func (service *mutableLargePayloadRunService) GetRunForAgent(ctx context.Context, workspaceID, runID string, _ workflowrun.AgentPrincipal) (workflowrunstore.Run, error) {
	return service.GetRun(ctx, workspaceID, runID)
}

func TestWorkflowPlanDraftSourceReadRejectsRevisionUpdateBetweenPages(t *testing.T) {
	firstSource := strings.Repeat("source-a\n", 6000)
	secondSource := strings.Repeat("source-b\n", 6000)
	snapshot := workflowpayload.Snapshot{Version: "7", Digest: workflowpayload.ContentDigest([]byte(firstSource))}
	first, err := workflowPayloadPageForTextWithSnapshot("plan", "wsp_test", "wpl_1", "", "", "", "source", "", "", 2048, firstSource, snapshot)
	require.NoError(t, err)
	require.NotEmpty(t, first.NextCursor)

	_, err = workflowPayloadPageForTextWithSnapshot("plan", "wsp_test", "wpl_1", "", "", "", "source", "", first.NextCursor, 2048, secondSource, workflowpayload.Snapshot{
		Version: "8", Digest: workflowpayload.ContentDigest([]byte(secondSource)),
	})
	require.ErrorIs(t, err, workflowpayload.ErrContentChanged)
}

func TestWorkflowRunSectionReadRejectsOutputUpdateBetweenPages(t *testing.T) {
	firstRaw := []byte(`{"payload":"` + strings.Repeat("a", 50000) + `"}`)
	secondRaw := []byte(`{"payload":"` + strings.Repeat("b", 50000) + `"}`)
	now := time.Now().UTC()
	service := &mutableLargePayloadRunService{run: workflowrunstore.Run{
		ID: "wri_1", WorkspaceID: "wsp_test", InvocationID: "wiv_1", PlanID: "wpl_1", PlanVersionID: "wpv_1",
		Status: workflowrunstore.RunStatusWaiting, InputJSON: []byte(`{}`), OutputJSON: firstRaw, StateJSON: []byte(`{}`), QueuedAt: now, CreatedAt: now, UpdatedAt: now,
	}}
	router := workflowPublicTestRouter(t, service, nil)
	response := postWorkflowPublicCommand(t, router, "/workflow/run", map[string]any{
		"action": "show", "runId": "wri_1", "section": "output", "limit": 2048,
	})
	require.Equal(t, 200, response.Code, response.Body.String())
	var envelope struct {
		Payload struct {
			NextCursor string `json:"nextCursor"`
		} `json:"payload"`
	}
	require.NoError(t, json.Unmarshal(response.Body.Bytes(), &envelope))
	require.NotEmpty(t, envelope.Payload.NextCursor)

	service.run.OutputJSON = secondRaw
	response = postWorkflowPublicCommand(t, router, "/workflow/run", map[string]any{
		"action": "show", "runId": "wri_1", "section": "output", "cursor": envelope.Payload.NextCursor, "limit": 2048,
	})
	require.Equal(t, 409, response.Code, response.Body.String())
	require.NotContains(t, response.Body.String(), string(secondRaw))
	require.Contains(t, response.Body.String(), "payload snapshot changed")
}

func TestWorkflowRunSectionReadRecoversOversizedInput(t *testing.T) {
	raw := []byte(`{"payload":"` + strings.Repeat("x", 50000) + `"}`)
	now := time.Now().UTC()
	router := workflowPublicTestRouter(t, largePayloadRunService{run: workflowrunstore.Run{
		ID: "wri_1", WorkspaceID: "wsp_test", InvocationID: "wiv_1", PlanID: "wpl_1", PlanVersionID: "wpv_1",
		Status: workflowrunstore.RunStatusWaiting, InputJSON: raw, StateJSON: []byte(`{}`), QueuedAt: now, CreatedAt: now, UpdatedAt: now,
	}}, nil)

	var cursor string
	var recovered strings.Builder
	for {
		response := postWorkflowPublicCommand(t, router, "/workflow/run", map[string]any{
			"action": "show", "runId": "wri_1", "section": "input", "cursor": cursor, "limit": 2048,
		})
		require.Equal(t, 200, response.Code, response.Body.String())
		var envelope struct {
			Payload struct {
				Value      string `json:"value"`
				NextCursor string `json:"nextCursor"`
			} `json:"payload"`
		}
		require.NoError(t, json.Unmarshal(response.Body.Bytes(), &envelope))
		recovered.WriteString(envelope.Payload.Value)
		if envelope.Payload.NextCursor == "" {
			break
		}
		cursor = envelope.Payload.NextCursor
	}
	require.Equal(t, string(raw), recovered.String())
}