package standard_test

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	agentcontext "github.com/yaoapp/yao/agent/context"
	"github.com/yaoapp/yao/agent/robot/executor/standard"
	executortypes "github.com/yaoapp/yao/agent/robot/executor/types"
	"github.com/yaoapp/yao/agent/robot/types"
	"github.com/yaoapp/yao/agent/testutils"
)

// ============================================================================
// RunExecution with ResumeContext tests
// ============================================================================

func TestRunExecutionResumeContext(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration test")
	}

	testutils.Prepare(t)
	defer testutils.Clean(t)

	ctx := types.NewContext(context.Background(), testAuth())

	t.Run("resumes from task index with previous results", func(t *testing.T) {
		robot := createRunTestRobot(t)
		exec := createRunTestExecution(robot)

		exec.Tasks = []types.Task{
			{
				ID:           "task-001",
				ExecutorType: types.ExecutorAssistant,
				ExecutorID:   "experts.text-writer",
				Messages: []agentcontext.Message{
					{Role: agentcontext.RoleUser, Content: "Write 'hello'"},
				},
				Order:  0,
				Status: types.TaskCompleted,
			},
			{
				ID:           "task-002",
				ExecutorType: types.ExecutorAssistant,
				ExecutorID:   "experts.text-writer",
				Messages: []agentcontext.Message{
					{Role: agentcontext.RoleUser, Content: "Write 'world'"},
				},
				Order:  1,
				Status: types.TaskPending,
			},
			{
				ID:           "task-003",
				ExecutorType: types.ExecutorAssistant,
				ExecutorID:   "experts.text-writer",
				Messages: []agentcontext.Message{
					{Role: agentcontext.RoleUser, Content: "Write '!'"},
				},
				Order:  2,
				Status: types.TaskPending,
			},
		}

		// Simulate resume: task-001 already completed, resume from task-002
		previousResult := types.TaskResult{
			TaskID:   "task-001",
			Success:  true,
			Output:   "hello",
			Duration: 100,
		}
		exec.ResumeContext = &types.ResumeContext{
			TaskIndex:       1,
			PreviousResults: []types.TaskResult{previousResult},
		}

		e := standard.New()
		err := e.RunExecution(ctx, exec, nil)

		require.NoError(t, err)
		// Should have 3 results: 1 from previous + 2 new
		require.Len(t, exec.Results, 3)

		assert.Equal(t, "task-001", exec.Results[0].TaskID)
		assert.True(t, exec.Results[0].Success)
		assert.Equal(t, "hello", exec.Results[0].Output)

		assert.Equal(t, "task-002", exec.Results[1].TaskID)
		assert.True(t, exec.Results[1].Success)

		assert.Equal(t, "task-003", exec.Results[2].TaskID)
		assert.True(t, exec.Results[2].Success)

		// ResumeContext should be cleared after completion
		assert.Nil(t, exec.ResumeContext)

		// Only task-002 and task-003 should have been executed (check status)
		assert.Equal(t, types.TaskCompleted, exec.Tasks[1].Status)
		assert.Equal(t, types.TaskCompleted, exec.Tasks[2].Status)
	})

	t.Run("resumes from last task", func(t *testing.T) {
		robot := createRunTestRobot(t)
		exec := createRunTestExecution(robot)

		exec.Tasks = []types.Task{
			{
				ID:           "task-001",
				ExecutorType: types.ExecutorAssistant,
				ExecutorID:   "experts.text-writer",
				Messages: []agentcontext.Message{
					{Role: agentcontext.RoleUser, Content: "Write 'hello'"},
				},
				Order:  0,
				Status: types.TaskCompleted,
			},
			{
				ID:           "task-002",
				ExecutorType: types.ExecutorAssistant,
				ExecutorID:   "experts.text-writer",
				Messages: []agentcontext.Message{
					{Role: agentcontext.RoleUser, Content: "Write 'world'"},
				},
				Order:  1,
				Status: types.TaskWaitingInput,
			},
		}

		// Resume from the last task
		exec.ResumeContext = &types.ResumeContext{
			TaskIndex: 1,
			PreviousResults: []types.TaskResult{
				{TaskID: "task-001", Success: true, Output: "hello", Duration: 100},
			},
		}

		e := standard.New()
		err := e.RunExecution(ctx, exec, nil)

		require.NoError(t, err)
		require.Len(t, exec.Results, 2)
		assert.True(t, exec.Results[1].Success)
		assert.Equal(t, types.TaskCompleted, exec.Tasks[1].Status)
	})
}

// ============================================================================
// Suspend method tests (using Executor directly)
// ============================================================================

func TestSuspendExecution(t *testing.T) {
	t.Run("suspend sets waiting fields and returns ErrExecutionSuspended", func(t *testing.T) {
		robot := &types.Robot{
			MemberID:    "test-robot-suspend",
			TeamID:      "test-team-1",
			DisplayName: "Suspend Test Robot",
		}

		exec := &types.Execution{
			ID:       "exec-suspend-001",
			MemberID: robot.MemberID,
			TeamID:   robot.TeamID,
			Status:   types.ExecRunning,
			Phase:    types.PhaseRun,
			Tasks: []types.Task{
				{ID: "task-001", Status: types.TaskRunning},
				{ID: "task-002", Status: types.TaskPending},
			},
			Results: []types.TaskResult{},
		}
		exec.SetRobot(robot)

		e := standard.NewWithConfig(executortypes.Config{SkipPersistence: true})
		err := e.Suspend(
			types.NewContext(context.Background(), nil),
			exec, 0, "What time range?",
		)

		assert.ErrorIs(t, err, types.ErrExecutionSuspended)
		assert.Equal(t, types.ExecWaiting, exec.Status)
		assert.Equal(t, "task-001", exec.WaitingTaskID)
		assert.Equal(t, "What time range?", exec.WaitingQuestion)
		assert.NotNil(t, exec.WaitingSince)
		assert.NotNil(t, exec.ResumeContext)
		assert.Equal(t, 0, exec.ResumeContext.TaskIndex)
		assert.Equal(t, types.TaskWaitingInput, exec.Tasks[0].Status)
	})

	t.Run("suspend with out of range taskIndex is safe", func(t *testing.T) {
		robot := &types.Robot{
			MemberID: "test-robot-suspend-2",
			TeamID:   "test-team-1",
		}
		exec := &types.Execution{
			ID:       "exec-suspend-002",
			MemberID: robot.MemberID,
			TeamID:   robot.TeamID,
			Status:   types.ExecRunning,
			Tasks:    []types.Task{},
			Results:  []types.TaskResult{},
		}
		exec.SetRobot(robot)

		e := standard.NewWithConfig(executortypes.Config{SkipPersistence: true})
		err := e.Suspend(
			types.NewContext(context.Background(), nil),
			exec, 5, "some question",
		)

		assert.ErrorIs(t, err, types.ErrExecutionSuspended)
		assert.Equal(t, types.ExecWaiting, exec.Status)
		assert.Empty(t, exec.WaitingTaskID)
	})
}

// ============================================================================
// ExecuteWithControl handles ErrExecutionSuspended
// ============================================================================

func TestExecuteWithControlSuspend(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration test")
	}

	testutils.Prepare(t)
	defer testutils.Clean(t)

	t.Run("returns ErrExecutionSuspended without marking as failed", func(t *testing.T) {
		// This test requires a robot-need-input assistant that returns need_input.
		// Since we don't have one yet (Stage 6), we test the suspend path indirectly
		// by verifying that when RunExecution returns ErrExecutionSuspended,
		// ExecuteWithControl propagates it correctly.
		//
		// Full E2E test with real assistant will be in Stage 6.

		robot := &types.Robot{
			MemberID:    "test-robot-suspend-exec",
			TeamID:      "test-team-1",
			DisplayName: "Suspend Exec Test",
			Config: &types.Config{
				Identity: &types.Identity{
					Role: "Test",
				},
				Resources: &types.Resources{
					Phases: map[types.Phase]string{},
					Agents: []string{"experts.text-writer"},
				},
				Quota: &types.Quota{Max: 5},
			},
		}

		ctx := types.NewContext(context.Background(), testAuth())
		e := standard.New()

		// Execute normally (no need_input expected from text-writer)
		exec, err := e.Execute(ctx, robot, types.TriggerHuman, "Write a greeting")
		if err == types.ErrExecutionSuspended {
			// If somehow suspended, verify state
			assert.Equal(t, types.ExecWaiting, exec.Status)
			assert.NotEmpty(t, exec.WaitingQuestion)
		} else {
			// Normal completion
			assert.NoError(t, err)
			assert.NotNil(t, exec)
		}
	})
}

// ============================================================================
// ResumeContext data structure tests
// ============================================================================

func TestResumeContext(t *testing.T) {
	t.Run("stores task index and previous results", func(t *testing.T) {
		rc := &types.ResumeContext{
			TaskIndex: 2,
			PreviousResults: []types.TaskResult{
				{TaskID: "t1", Success: true, Output: "out1"},
				{TaskID: "t2", Success: false, Error: "some error"},
			},
		}
		assert.Equal(t, 2, rc.TaskIndex)
		assert.Len(t, rc.PreviousResults, 2)
		assert.True(t, rc.PreviousResults[0].Success)
		assert.False(t, rc.PreviousResults[1].Success)
	})
}

// ============================================================================
// NeedInput in TaskResult
// ============================================================================

func TestTaskResultNeedInput(t *testing.T) {
	t.Run("NeedInput fields are populated correctly", func(t *testing.T) {
		result := types.TaskResult{
			TaskID:        "task-001",
			Success:       true,
			Output:        "some output",
			NeedInput:     true,
			InputQuestion: "What time range?",
		}
		assert.True(t, result.NeedInput)
		assert.Equal(t, "What time range?", result.InputQuestion)
	})
}

// ============================================================================
// Execution status transitions for suspend/resume
// ============================================================================

func TestExecutionStatusTransitions(t *testing.T) {
	t.Run("ExecWaiting is a valid status", func(t *testing.T) {
		exec := &types.Execution{
			ID:     "exec-001",
			Status: types.ExecWaiting,
		}
		assert.Equal(t, types.ExecStatus("waiting"), exec.Status)
	})

	t.Run("TaskWaitingInput is a valid task status", func(t *testing.T) {
		task := types.Task{
			ID:     "task-001",
			Status: types.TaskWaitingInput,
		}
		assert.Equal(t, types.TaskStatus("waiting_input"), task.Status)
	})

	t.Run("Execution V2 fields are accessible", func(t *testing.T) {
		now := time.Now()
		exec := &types.Execution{
			ID:              "exec-v2-001",
			ChatID:          "robot_member1_exec001",
			WaitingTaskID:   "task-002",
			WaitingQuestion: "What period?",
			WaitingSince:    &now,
			ResumeContext: &types.ResumeContext{
				TaskIndex: 1,
				PreviousResults: []types.TaskResult{
					{TaskID: "task-001", Success: true},
				},
			},
		}
		assert.Equal(t, "robot_member1_exec001", exec.ChatID)
		assert.Equal(t, "task-002", exec.WaitingTaskID)
		assert.Equal(t, "What period?", exec.WaitingQuestion)
		assert.NotNil(t, exec.WaitingSince)
		assert.NotNil(t, exec.ResumeContext)
		assert.Equal(t, 1, exec.ResumeContext.TaskIndex)
	})
}