package run

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

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/workflow/plan"
)

const countedAgentTeamRuntimePlan = `apiVersion: loop.plan/v1
kind: PlanVersion
input:
  answer: {type: string, required: true}
output:
  answer: {type: string, required: true}
  iteration: {type: number, required: true}
  last: {type: boolean, required: true}
graph:
  entry:
    nodeId: repeat
    port: enter
    input:
      answer: {ref: plan.input.answer}
      iteration: {value: 0}
      last: {value: false}
  nodes:
  - id: repeat
    type: loop
    loop:
      mode: counted
      iterations: 3
      state:
        answer: {type: string, required: true}
        iteration: {type: number, required: true}
        last: {type: boolean, required: true}
  - id: success
    type: end
    end: {outcome: succeeded}
  - id: failure
    type: end
    end: {outcome: failed, code: loop_failed}
  edges:
  - id: body
    source: {nodeId: repeat, port: body}
    target: {nodeId: repeat, port: continue}
    bindings:
      answer: {ref: loop.state.answer}
      iteration: {ref: loop.iteration}
      last: {ref: loop.isLast}
  - id: done
    source: {nodeId: repeat, port: done}
    target: {nodeId: success, port: in}
  - id: failed
    source: {nodeId: repeat, port: failed}
    target: {nodeId: failure, port: in}
`

const countedLoopEarlyBreakPlan = `apiVersion: loop.plan/v1
kind: PlanVersion
input:
  answer: {type: string, required: true}
  stop: {type: boolean, required: true}
output:
  answer: {type: string, required: true}
  stop: {type: boolean, required: true}
graph:
  entry:
    nodeId: repeat
    port: enter
    input:
      answer: {ref: plan.input.answer}
      stop: {ref: plan.input.stop}
  nodes:
  - id: repeat
    type: loop
    loop:
      mode: counted
      iterations: 5
      state:
        answer: {type: string, required: true}
        stop: {type: boolean, required: true}
  - id: route
    type: switch
    input:
      answer: {type: string, required: true}
      stop: {type: boolean, required: true}
    switch:
      mode: first_match
      cases:
      - port: stop
        when: input.stop == true
      - port: again
        isDefault: true
  - id: success
    type: end
    end: {outcome: succeeded}
  - id: failure
    type: end
    end: {outcome: failed, code: loop_failed}
  edges:
  - id: body
    source: {nodeId: repeat, port: body}
    target: {nodeId: route, port: in}
  - id: stop
    source: {nodeId: route, port: stop}
    target: {nodeId: repeat, port: break}
  - id: again
    source: {nodeId: route, port: again}
    target: {nodeId: repeat, port: continue}
  - id: done
    source: {nodeId: repeat, port: done}
    target: {nodeId: success, port: in}
  - id: failed
    source: {nodeId: repeat, port: failed}
    target: {nodeId: failure, port: in}
`

func TestCountedLoopRunsExactlyNBodyTraversalsThenAutomaticallySucceeds(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, countedAgentTeamRuntimePlan)
	state, err := kernel.Start(Values{"answer": "开门见山"})
	require.NoError(t, err)

	events := []Event{}
	for step := 0; step < 10 && !state.Status.Terminal(); step++ {
		transition, stepErr := kernel.Step(state, reducerNow)
		require.NoError(t, stepErr)
		state = transition.State
		events = append(events, transition.Events...)
		if step == 0 {
			require.Len(t, state.Scopes, 1)
			require.Equal(t, plan.LoopCounted, state.Scopes[0].Loop.Mode)
			require.Equal(t, 3, state.Scopes[0].Loop.Iterations)
			restored, roundTripErr := JSONRoundTrip(state)
			require.NoError(t, roundTripErr)
			require.Equal(t, state, restored)
		}
	}

	require.Equal(t, StatusSucceeded, state.Status)
	require.Equal(t, "开门见山", state.Output["answer"])
	require.Equal(t, json.Number("3"), state.Output["iteration"])
	require.Equal(t, true, state.Output["last"])
	require.Len(t, state.Scopes, 1)
	require.Equal(t, ScopeSucceeded, state.Scopes[0].Status)
	require.Equal(t, 3, state.Scopes[0].Loop.Iteration)

	body, done := 0, 0
	for _, activation := range state.Activations {
		if activation.NodeID != "repeat" {
			continue
		}
		for _, port := range activation.SelectedPorts {
			switch port {
			case "body":
				body++
			case "done":
				done++
			case "failed":
				t.Fatalf("counted loop unexpectedly failed: %#v", activation)
			}
		}
	}
	require.Equal(t, 3, body, "counted loop must expose body exactly N times")
	require.Equal(t, 1, done)
	require.Equal(t, 1, countEventKind(events, "loop_opened"))
	require.Equal(t, 2, countEventKind(events, "loop_continued"))
	require.Equal(t, 1, countEventKind(events, "loop_completed"))
}

func TestCountedLoopSupportsOptionalEarlyBreak(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, countedLoopEarlyBreakPlan)
	state, err := kernel.Start(Values{"answer": "开门见山", "stop": true})
	require.NoError(t, err)
	terminal := stepUntilTerminal(t, kernel, state, 8)
	require.Equal(t, StatusSucceeded, terminal.Status)
	require.Equal(t, "开门见山", terminal.Output["answer"])
	require.Equal(t, true, terminal.Output["stop"])
	require.Len(t, terminal.Scopes, 1)
	require.Equal(t, 1, terminal.Scopes[0].Loop.Iteration)
	require.Equal(t, ScopeSucceeded, terminal.Scopes[0].Status)

	body := 0
	for _, activation := range terminal.Activations {
		if activation.NodeID == "repeat" && len(activation.SelectedPorts) == 1 && activation.SelectedPorts[0] == "body" {
			body++
		}
	}
	require.Equal(t, 1, body)
}

func TestCountedLoopOneIterationIsNotOffByOne(t *testing.T) {
	t.Parallel()
	source := strings.Replace(countedAgentTeamRuntimePlan, "iterations: 3", "iterations: 1", 1)
	kernel := compilePlanKernel(t, source)
	state, err := kernel.Start(Values{"answer": "一心一意"})
	require.NoError(t, err)
	terminal := stepUntilTerminal(t, kernel, state, 4)
	require.Equal(t, StatusSucceeded, terminal.Status)
	require.Equal(t, json.Number("1"), terminal.Output["iteration"])
	require.Equal(t, true, terminal.Output["last"])

	body := 0
	for _, activation := range terminal.Activations {
		if activation.NodeID == "repeat" && len(activation.SelectedPorts) == 1 && activation.SelectedPorts[0] == "body" {
			body++
		}
	}
	require.Equal(t, 1, body)
}

func compilePlanKernel(t *testing.T, source string) *Kernel {
	t.Helper()
	artifact, diagnostics, err := plan.Compile([]byte(source))
	require.NoError(t, err)
	require.False(t, plan.HasErrors(diagnostics), "diagnostics: %#v", diagnostics)
	wire, err := artifact.DeterministicJSON()
	require.NoError(t, err)
	graph, err := DecodeArtifact(wire)
	require.NoError(t, err)
	kernel, err := NewKernel(graph)
	require.NoError(t, err)
	return kernel
}

func countEventKind(events []Event, kind string) int {
	count := 0
	for _, event := range events {
		if event.Kind == kind {
			count++
		}
	}
	return count
}