package agent_test

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api/agent"
	"github.com/tidbcloud/AgentTeam/internal/api/dto"
	"github.com/tidbcloud/AgentTeam/internal/api/testkit"
	"github.com/tidbcloud/AgentTeam/internal/archive"
	"github.com/tidbcloud/AgentTeam/internal/core/secrets"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
	"github.com/tidbcloud/AgentTeam/internal/store/workflow/plan"
	"github.com/tidbcloud/AgentTeam/internal/workflow/planmanage"
)

func TestWorkflowReplyAcceptsExactlyOneTerminalMessageBeforeRuntimeReceipt(t *testing.T) {
	ctx := context.Background()
	_, pool := startAgentAPITestPool(t, ctx)
	archiveStorage, err := archive.NewFilesystem(t.TempDir())
	require.NoError(t, err)
	dataStore := testkit.New(pool, archiveStorage, 30)
	taskStore := mustTaskClient(t, pool)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
	dependencies := agent.New(
		logger,
		dataStore.AgentStorePortsDeps(),
		dataStore.Attachment(),
		taskStore,
		nil,
		nil,
	)
	dependencies.WorkflowReplies = planmanage.NewAgentEditService(
		pool,
		plan.New(pool),
		dataStore.Message(),
		dataStore.App(),
		dataStore.AgentOps(),
	)
	agent.RegisterRoutes(router, dependencies)

	fixture := seedWorkflowReplyAPIFixture(t, pool)
	path := "/api/v1/workspaces/" + fixture.workspaceID +
		"/agents/" + fixture.agentID + "/runtime/send"
	payload := func(clientMessageID, content string, seenUpToSeq int64, sourceMessageID string) []byte {
		raw, marshalErr := json.Marshal(map[string]any{
			"target":          "workflow:" + fixture.planID,
			"content":         content,
			"clientMessageId": clientMessageID,
			"freshness": map[string]any{
				"targetKey":        fixture.scopeID,
				"seenUpToSeq":      seenUpToSeq,
				"sourceMessageIds": []string{sourceMessageID},
			},
		})
		require.NoError(t, marshalErr)
		return raw
	}
	otherPath := "/api/v1/workspaces/" + fixture.workspaceID +
		"/agents/" + fixture.otherAgentID + "/runtime/send"
	nonSelected := doRawAgentRequest(
		router,
		http.MethodPost,
		fixture.otherAgentID,
		fixture.otherAgentToken,
		otherPath,
		payload("client_workflow_reply_other", "Membership must not authorize this reply.", 1, fixture.requestMessageID),
	)
	require.Equal(t, http.StatusForbidden, nonSelected.Code, nonSelected.Body.String())

	_, err = pool.Exec(ctx, `
DELETE FROM scope_fanout_agent_deliveries
WHERE source_event_type = 'message'
  AND source_event_id = $1
  AND agent_id = $2
`, fixture.requestMessageID, fixture.agentID)
	require.NoError(t, err)
	withoutReceipt := doRawAgentRequest(
		router,
		http.MethodPost,
		fixture.agentID,
		fixture.agentToken,
		path,
		payload("client_workflow_reply_without_receipt", "A missing delivery receipt must reject this reply.", 1, fixture.requestMessageID),
	)
	require.Equal(t, http.StatusForbidden, withoutReceipt.Code, withoutReceipt.Body.String())
	_, err = pool.Exec(ctx, `
INSERT INTO scope_fanout_agent_deliveries (
  source_event_type, source_event_id, agent_id, outbox_id, delivered_at
) VALUES ('message', $1, $2, $3, $4)
`, fixture.requestMessageID, fixture.agentID, "sfo_workflow_reply_once", time.Now().UTC())
	require.NoError(t, err)

	oldSource := doRawAgentRequest(
		router,
		http.MethodPost,
		fixture.agentID,
		fixture.agentToken,
		path,
		payload("client_workflow_reply_old_source", "The current request source is required.", 1, "msg_workflow_reply_old_request"),
	)
	require.Equal(t, http.StatusForbidden, oldSource.Code, oldSource.Body.String())

	runtimeBase := "/api/v1/workspaces/" + fixture.workspaceID +
		"/agents/" + fixture.agentID + "/runtime"
	history := doRawAgentRequest(
		router,
		http.MethodGet,
		fixture.agentID,
		fixture.agentToken,
		runtimeBase+"/history?scope_id="+fixture.scopeID,
		nil,
	)
	require.Equal(t, http.StatusOK, history.Code, history.Body.String())
	var historyBody dto.AgentHistoryResponse
	require.NoError(t, json.Unmarshal(history.Body.Bytes(), &historyBody))
	require.Empty(t, historyBody.Messages)

	readCheck := doRawAgentRequest(
		router,
		http.MethodGet,
		fixture.agentID,
		fixture.agentToken,
		runtimeBase+"/messages/"+fixture.requestMessageID+"/read-check",
		nil,
	)
	require.Equal(t, http.StatusOK, readCheck.Code, readCheck.Body.String())
	var readCheckBody dto.AgentReadCheckResponse
	require.NoError(t, json.Unmarshal(readCheck.Body.Bytes(), &readCheckBody))
	require.False(t, readCheckBody.Allowed)

	memberList := doRawAgentRequest(
		router,
		http.MethodGet,
		fixture.agentID,
		fixture.agentToken,
		runtimeBase+"/scopes/"+fixture.scopeID+"/members",
		nil,
	)
	require.Equal(t, http.StatusOK, memberList.Code, memberList.Body.String())
	var memberListBody dto.AgentScopeMembersResponse
	require.NoError(t, json.Unmarshal(memberList.Body.Bytes(), &memberListBody))
	require.Empty(t, memberListBody.Members)

	newerContextDigest, err := message.ComputeMessageCreatePayloadDigest(message.MessageCreateDigestInput{
		MessageType: "text",
		Content:     "Newer Workflow context.",
	})
	require.NoError(t, err)
	require.NoError(t, coredb.WithTx(ctx, pool, func(tx coredb.Tx) error {
		_, writeErr := message.WriteWorkflowScopeMessageTx(ctx, tx, message.WorkflowScopeMessageWrite{
			WorkspaceID:     fixture.workspaceID,
			ScopeID:         fixture.scopeID,
			ChannelName:     fixture.planID,
			AuthorMemberID:  fixture.otherAgentMemberID,
			AuthorType:      message.AuthorActorTypeAgent,
			ClientMessageID: "client_workflow_reply_newer_context",
			PayloadDigest:   newerContextDigest,
			Content:         "Newer Workflow context.",
			AuthorProfile: message.AuthorProfile{
				ActorType:   message.AuthorActorTypeAgent,
				PrincipalID: fixture.otherAgentID,
				MemberID:    fixture.otherAgentMemberID,
				TargetID:    "other-workflow-editor",
				DisplayName: "Other Workflow Editor",
			},
		})
		return writeErr
	}))
	held := doRawAgentRequest(
		router,
		http.MethodPost,
		fixture.agentID,
		fixture.agentToken,
		path,
		payload("client_workflow_reply_held", "Refresh before accepting this reply.", 1, fixture.requestMessageID),
	)
	require.Equal(t, http.StatusOK, held.Code, held.Body.String())
	var heldBody dto.AgentSendMessageResponse
	require.NoError(t, json.Unmarshal(held.Body.Bytes(), &heldBody))
	require.Equal(t, message.SendStateHeld, heldBody.State)
	require.Equal(t, int64(2), heldBody.NeedsRefreshUntilSeq)

	requests := []struct {
		clientMessageID string
		body            []byte
	}{
		{
			clientMessageID: "client_workflow_reply_first",
			body:            payload("client_workflow_reply_first", "First terminal reply.", 2, fixture.requestMessageID),
		},
		{
			clientMessageID: "client_workflow_reply_second",
			body:            payload("client_workflow_reply_second", "Second terminal reply.", 2, fixture.requestMessageID),
		},
	}

	start := make(chan struct{})
	responses := make(chan workflowReplyAttempt, len(requests))
	for _, request := range requests {
		request := request
		go func() {
			<-start
			httpRequest := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(request.body))
			httpRequest.Header.Set("Content-Type", "application/json")
			httpRequest.Header.Set("Authorization", "Bearer "+fixture.agentToken)
			httpRequest.Header.Set("X-Agent-Id", fixture.agentID)
			response := httptest.NewRecorder()
			router.ServeHTTP(response, httpRequest)
			responses <- workflowReplyAttempt{
				clientMessageID: request.clientMessageID,
				response:        response,
			}
		}()
	}
	close(start)

	var winner workflowReplyAttempt
	statusCounts := map[int]int{}
	responseBodies := map[int][]string{}
	for range requests {
		attempt := <-responses
		statusCounts[attempt.response.Code]++
		responseBodies[attempt.response.Code] = append(
			responseBodies[attempt.response.Code],
			attempt.response.Body.String(),
		)
		if attempt.response.Code == http.StatusOK {
			winner = attempt
		}
	}
	require.Equal(t, 1, statusCounts[http.StatusOK], responseBodies)
	require.Equal(t, 1, statusCounts[http.StatusForbidden], responseBodies)
	require.NotEmpty(t, winner.clientMessageID)

	var sent dto.AgentSendMessageResponse
	require.NoError(t, json.Unmarshal(winner.response.Body.Bytes(), &sent))
	require.Equal(t, message.SendStateSent, sent.State)
	require.NotEmpty(t, sent.MessageID)

	var messageCount int
	require.NoError(t, pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM scope_messages
WHERE scope_id = $1
`, fixture.scopeID).Scan(&messageCount))
	require.Equal(t, 3, messageCount)

	var terminalEventCount int
	require.NoError(t, pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM agent_runtime_events
WHERE workspace_id = $1
  AND agent_id = $2
  AND message_id = $3
  AND event_class = 'visible_reply_sent'
  AND reply_message_id = $4
`, fixture.workspaceID, fixture.agentID, fixture.requestMessageID, sent.MessageID).Scan(&terminalEventCount))
	require.Equal(t, 1, terminalEventCount, "the final reply commit must also close the wake durably")

	var winnerBody []byte
	for _, request := range requests {
		if request.clientMessageID == winner.clientMessageID {
			winnerBody = request.body
			break
		}
	}
	replay := doRawAgentRequest(
		router, http.MethodPost, fixture.agentID, fixture.agentToken, path, winnerBody,
	)
	require.Equal(t, http.StatusOK, replay.Code, replay.Body.String())
	var replayed dto.AgentSendMessageResponse
	require.NoError(t, json.Unmarshal(replay.Body.Bytes(), &replayed))
	require.Equal(t, message.SendStateAlreadySent, replayed.State)
	require.Equal(t, sent.MessageID, replayed.MessageID)

	changedReplay := doRawAgentRequest(
		router,
		http.MethodPost,
		fixture.agentID,
		fixture.agentToken,
		path,
		payload(winner.clientMessageID, "A reused client id cannot change payload.", 2, fixture.requestMessageID),
	)
	require.Equal(t, http.StatusConflict, changedReplay.Code, changedReplay.Body.String())

	afterTerminal := doRawAgentRequest(
		router,
		http.MethodPost,
		fixture.agentID,
		fixture.agentToken,
		path,
		payload("client_workflow_reply_after_terminal", "A terminal request cannot accept another reply.", 2, fixture.requestMessageID),
	)
	require.Equal(t, http.StatusForbidden, afterTerminal.Code, afterTerminal.Body.String())
}

type workflowReplyAttempt struct {
	clientMessageID string
	response        *httptest.ResponseRecorder
}

type workflowReplyAPIFixture struct {
	workspaceID        string
	planID             string
	scopeID            string
	agentID            string
	agentToken         string
	otherAgentID       string
	otherAgentToken    string
	otherAgentMemberID string
	requestMessageID   string
}

func seedWorkflowReplyAPIFixture(
	t *testing.T,
	pool coredb.QueryExecer,
) workflowReplyAPIFixture {
	t.Helper()
	ctx := context.Background()
	now := time.Date(2026, 8, 7, 12, 0, 0, 0, time.UTC)
	fixture := workflowReplyAPIFixture{
		workspaceID:        "wsp_workflow_reply_once",
		planID:             "pln_workflow_reply_once",
		scopeID:            "wfs_workflow_reply_once",
		agentID:            "agt_000000workflow_reply_once",
		agentToken:         "workflow-reply-once-token",
		otherAgentID:       "agt_000000workflow_reply_once_other",
		otherAgentToken:    "workflow-reply-once-other-token",
		otherAgentMemberID: "9pmp13zniqpm",
		requestMessageID:   "msg_workflow_reply_once_request",
	}
	const (
		userID         = "usr_000000workflow_reply_once"
		humanMemberID  = "9ra2mp9eadi0"
		agentMemberID  = "g7mv6bhnhlf8"
		machineID      = "mac_workflow_reply_once"
		otherMachineID = "mac_workflow_reply_once_other"
	)
	mustExec(t, pool, `
INSERT INTO users (
  id, email, name, display_name, avatar_url, is_admin,
  password_hash, email_verified_at, created_at, updated_at
) VALUES (
  $1, 'workflow-reply-once@example.com', 'workflow-reply-once',
  'Workflow Reply Owner', '', FALSE, 'hash', $2, $2, $2
)
`, userID, now)
	organizationID := seedAgentTestWorkspaceOrganization(
		t, pool, fixture.workspaceID, "workflow-reply-once", userID, now,
	)
	mustExec(t, pool, `
INSERT INTO workspaces (
  id, organization_id, slug, name, latest_message_seq, created_at, updated_at
) VALUES ($1, $2, 'workflow-reply-once', 'Workflow Reply Once', 1, $3, $3)
`, fixture.workspaceID, organizationID, now)
	mustExec(t, pool, `
INSERT INTO members (
  id, workspace_id, type, display_name, user_id, workspace_role,
  status, created_at, updated_at
) VALUES ($1, $2, 'human', 'Owner', $3, 'owner', 'active', $4, $4)
`, humanMemberID, fixture.workspaceID, userID, now)
	mustExec(t, pool, `
INSERT INTO members (
  id, workspace_id, type, display_name, status, created_at, updated_at
) VALUES ($1, $2, 'agent', 'Workflow Editor', 'active', $3, $3)
`, agentMemberID, fixture.workspaceID, now)
	mustExec(t, pool, `
INSERT INTO members (
  id, workspace_id, type, display_name, status, created_at, updated_at
) VALUES ($1, $2, 'agent', 'Other Workflow Editor', 'active', $3, $3)
`, fixture.otherAgentMemberID, fixture.workspaceID, now)
	mustExec(t, pool, `
INSERT INTO machines (
  id, workspace_id, created_by_user_id, name, api_key, status, created_at, updated_at
) VALUES ($1, $2, $3, 'Workflow Reply Machine', 'machine-key', 'active', $4, $4)
`, machineID, fixture.workspaceID, userID, now)
	mustExec(t, pool, `
INSERT INTO machines (
  id, workspace_id, created_by_user_id, name, api_key, status, created_at, updated_at
) VALUES ($1, $2, $3, 'Other Workflow Reply Machine', 'other-machine-key', 'active', $4, $4)
`, otherMachineID, fixture.workspaceID, userID, now)
	mustExec(t, pool, `
INSERT INTO agents (
  id, workspace_id, machine_id, member_id, name, runtime, model,
  description, auth_token_hash, status, activity_detail, created_at, updated_at
) VALUES (
  $1, $2, $3, $4, 'workflow-editor', 'codex', 'gpt-5',
  '', $5, 'active', '', $6, $6
)
`, fixture.agentID, fixture.workspaceID, machineID, agentMemberID, secrets.HashToken(fixture.agentToken), now)
	mustExec(t, pool, `
INSERT INTO agents (
  id, workspace_id, machine_id, member_id, name, runtime, model,
  description, auth_token_hash, status, activity_detail, created_at, updated_at
) VALUES (
  $1, $2, $3, $4, 'other-workflow-editor', 'codex', 'gpt-5',
  '', $5, 'active', '', $6, $6
)
`, fixture.otherAgentID, fixture.workspaceID, otherMachineID, fixture.otherAgentMemberID, secrets.HashToken(fixture.otherAgentToken), now)

	repositories := message.NewScopeRepositories(pool)
	require.NoError(t, repositories.Scopes.UpsertScope(ctx, message.ScopeRecord{
		ID: fixture.scopeID, WorkspaceID: fixture.workspaceID,
		Type: message.WorkflowScopeType, IdentityKey: "workflow:" + fixture.planID,
		LatestSeq: 1, LatestGlobalSeq: 1, Status: "active",
		CreatedByMemberID: humanMemberID, CreatedAt: now, UpdatedAt: now,
	}))
	for _, memberID := range []string{humanMemberID, agentMemberID, fixture.otherAgentMemberID} {
		require.NoError(t, repositories.Memberships.Upsert(ctx, message.ScopeMembershipRecord{
			ScopeID: fixture.scopeID, MemberID: memberID, Visibility: "public",
			EffectiveFromSeq: 1, CreatedAt: now, UpdatedAt: now,
		}))
	}
	mustExec(t, pool, `
INSERT INTO workflow_plans (
  id, workspace_id, scope_id, plan_key, display_name, description_md, revision,
  create_idempotency_key, create_request_digest,
  created_by_member_id, created_by_actor_type, created_by_actor_id,
  updated_by_member_id, updated_by_actor_type, updated_by_actor_id,
  created_at, updated_at
) VALUES (
  $1, $2, $3, $1, 'Workflow Reply Once', '', 1,
  $1, 'digest',
  $4, 'human', $5,
  $4, 'human', $5,
  $6, $6
)
`, fixture.planID, fixture.workspaceID, fixture.scopeID, humanMemberID, userID, now)
	metadata, _, err := message.EncodeMessageMetadata(map[string]any{
		message.WorkflowAgentEditMetadataKey: map[string]any{
			"kind":              "request",
			"planId":            fixture.planID,
			"baseDraftRevision": int64(1),
			"responderAgentId":  fixture.agentID,
			"target": map[string]any{
				"kind": "node", "id": "notify", "nodeType": "send_message",
			},
		},
	})
	require.NoError(t, err)
	require.NoError(t, repositories.Messages.Insert(ctx, message.ScopeMessageRecord{
		ID: fixture.requestMessageID, WorkspaceID: fixture.workspaceID,
		ScopeID: fixture.scopeID, Seq: 1, WorkspaceGlobalSeq: 1,
		AuthorMemberID: humanMemberID, AuthorType: message.AuthorActorTypeHuman,
		MessageType: "text", Content: "Update the Workflow.",
		Metadata: metadata, CreatedAt: now, UpdatedAt: now,
	}))
	mustExec(t, pool, `
INSERT INTO scope_fanout_outbox (
  id, workspace_id, scope_id, message_id,
  event_ref_type, event_ref_id, fanout_type, dedupe_key, payload,
  status, attempts, max_attempts, next_attempt_at,
  created_at, delivered_at, last_error
) VALUES (
  $1, $2, $3, $4,
  'message', $4, 'scope_fanout', $1, JSON_OBJECT(),
  'delivered', 1, 10, $5,
  $5, $5, ''
)
`, "sfo_workflow_reply_once", fixture.workspaceID, fixture.scopeID, fixture.requestMessageID, now)
	mustExec(t, pool, `
INSERT INTO scope_fanout_agent_deliveries (
  source_event_type, source_event_id, agent_id, outbox_id, delivered_at
) VALUES ('message', $1, $2, $3, $4)
`, fixture.requestMessageID, fixture.agentID, "sfo_workflow_reply_once", now)
	return fixture
}

func doRawAgentRequest(
	router http.Handler,
	method string,
	agentID, token, path string,
	body []byte,
) *httptest.ResponseRecorder {
	request := httptest.NewRequest(method, path, bytes.NewReader(body))
	if len(body) != 0 {
		request.Header.Set("Content-Type", "application/json")
	}
	request.Header.Set("Authorization", "Bearer "+token)
	request.Header.Set("X-Agent-Id", agentID)
	response := httptest.NewRecorder()
	router.ServeHTTP(response, request)
	return response
}