package service

import (
	"context"
	"fmt"
	"strings"
	"testing"

	"ai-agent-scaffold-go/internal/model"
	"github.com/stretchr/testify/assert"
)

// ============================================================
// Stub 实现
// ============================================================

// stubChatModel 实现 ChatModelWithTools 接口
type stubChatModel struct {
	generateReply model.ChatReply
	generateErr   error
	streamDelta   string
	streamErr     error
	toolResult    string
	toolErr       error
}

func (m *stubChatModel) Generate(ctx context.Context, msgs []model.ChatMessage) (model.ChatReply, error) {
	return m.generateReply, m.generateErr
}

func (m *stubChatModel) Stream(ctx context.Context, msgs []model.ChatMessage) (<-chan model.ChatStreamEvent, <-chan error) {
	events := make(chan model.ChatStreamEvent, 4)
	errs := make(chan error, 1)

	go func() {
		defer close(events)
		defer close(errs)
		if m.streamErr != nil {
			errs <- m.streamErr
			return
		}
		// 模拟逐字输出
		for _, ch := range m.streamDelta {
			events <- model.ChatStreamEvent{Delta: string(ch)}
		}
		events <- model.ChatStreamEvent{Done: true}
	}()

	return events, errs
}

func (m *stubChatModel) CallTool(ctx context.Context, name, arguments string) (string, error) {
	return m.toolResult, m.toolErr
}

// stubToolCallModel 每次都返回工具调用的 stub
type stubToolCallModel struct {
	callCount int
	maxCalls  int // 达到此次数后返回文本
}

func (m *stubToolCallModel) Generate(ctx context.Context, msgs []model.ChatMessage) (model.ChatReply, error) {
	m.callCount++
	if m.callCount > m.maxCalls {
		return model.ChatReply{Content: "done"}, nil
	}
	return model.ChatReply{ToolCalls: []model.ChatToolCall{
		{ID: fmt.Sprintf("call_%d", m.callCount), Name: "tool1", Arguments: `{"query":"test"}`},
	}}, nil
}

func (m *stubToolCallModel) Stream(ctx context.Context, msgs []model.ChatMessage) (<-chan model.ChatStreamEvent, <-chan error) {
	events := make(chan model.ChatStreamEvent, 4)
	errs := make(chan error, 1)
	m.callCount++
	if m.callCount > m.maxCalls {
		events <- model.ChatStreamEvent{Delta: "done", Done: true}
	} else {
		events <- model.ChatStreamEvent{ToolCalls: []model.ChatToolCall{
			{ID: fmt.Sprintf("call_%d", m.callCount), Name: "tool1", Arguments: `{"query":"test"}`},
		}}
	}
	close(events)
	close(errs)
	return events, errs
}

func (m *stubToolCallModel) CallTool(ctx context.Context, name, arguments string) (string, error) {
	return "tool-result", nil
}

// ============================================================
// 辅助函数
// ============================================================

func newTestLLMAgent(cm ChatModelWithTools) *LLMAgent {
	return NewLLMAgent("test-agent", "you are a test agent", "test desc", "", cm)
}

func testContent(msg string) model.ChatContent {
	return model.ChatContent{Texts: []model.TextPart{{Message: msg}}}
}

// ============================================================
// LLMAgent 测试
// ============================================================

func TestLLMAgent_Run_ReturnsContent(t *testing.T) {
	cm := &stubChatModel{generateReply: model.ChatReply{Content: "hello"}}
	agent := newTestLLMAgent(cm)

	result, err := agent.Run(context.Background(), testContent("hi"))
	assert.NoError(t, err)
	assert.Equal(t, "hello", result)
}

func TestLLMAgent_Run_ToolCallLoop_TwoRounds(t *testing.T) {
	cm := &stubToolCallModel{maxCalls: 1}
	agent := newTestLLMAgent(cm)

	result, err := agent.Run(context.Background(), testContent("hi"))
	assert.NoError(t, err)
	assert.Equal(t, "done", result)
	assert.Equal(t, 2, cm.callCount)
}

func TestLLMAgent_Run_ExceedsIterationLimit_ReturnsError(t *testing.T) {
	cm := &stubToolCallModel{maxCalls: 100} // 永远返回工具调用
	agent := newTestLLMAgent(cm)

	_, err := agent.Run(context.Background(), testContent("hi"))
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "exceeded tool-call iteration limit")
}

func TestLLMAgent_Run_GenerateError_ReturnsError(t *testing.T) {
	cm := &stubChatModel{generateErr: fmt.Errorf("llm down")}
	agent := newTestLLMAgent(cm)

	_, err := agent.Run(context.Background(), testContent("hi"))
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "llm down")
}

func TestLLMAgent_Run_ToolCallError_ReturnsError(t *testing.T) {
	cm := &stubChatModel{
		generateReply: model.ChatReply{ToolCalls: []model.ChatToolCall{
			{ID: "c1", Name: "bad-tool", Arguments: "{}"},
		}},
		toolErr: fmt.Errorf("tool failed"),
	}
	agent := newTestLLMAgent(cm)

	_, err := agent.Run(context.Background(), testContent("hi"))
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "tool failed")
}

func TestLLMAgent_Stream_ReturnsChunks(t *testing.T) {
	cm := &stubChatModel{streamDelta: "hello"}
	agent := newTestLLMAgent(cm)

	out := make(chan string, 10)
	err := agent.Stream(context.Background(), testContent("hi"), out)
	assert.NoError(t, err)
	close(out)

	var texts []string
	for s := range out {
		texts = append(texts, s)
	}
	assert.Equal(t, []string{"h", "e", "l", "l", "o"}, texts)
}

func TestLLMAgent_Stream_Error_ReturnsError(t *testing.T) {
	cm := &stubChatModel{streamErr: fmt.Errorf("stream failed")}
	agent := newTestLLMAgent(cm)

	out := make(chan string, 10)
	err := agent.Stream(context.Background(), testContent("hi"), out)
	assert.Error(t, err)
}

func TestLLMAgent_Name_ReturnsName(t *testing.T) {
	cm := &stubChatModel{}
	agent := NewLLMAgent("my-agent", "inst", "desc", "key", cm)
	assert.Equal(t, "my-agent", agent.Name())
	assert.Equal(t, "key", agent.OutputKey())
}

// ============================================================
// SequentialAgent 测试
// ============================================================

func TestSequentialAgent_Run_ExecutesInOrder(t *testing.T) {
	cm1 := &stubChatModel{generateReply: model.ChatReply{Content: "first"}}
	cm2 := &stubChatModel{generateReply: model.ChatReply{Content: "second"}}

	sub1 := NewLLMAgent("a1", "inst1", "", "out1", cm1)
	sub2 := NewLLMAgent("a2", "inst2", "", "", cm2)

	seq := NewSequentialAgent("seq", "desc", []model.Agent{sub1, sub2})
	result, err := seq.Run(context.Background(), testContent("hi"))

	assert.NoError(t, err)
	assert.Equal(t, "second", result) // 返回最后一个的结果
}

func TestSequentialAgent_Run_PassesOutputKey(t *testing.T) {
	// sub1 输出 "first"，存入 vars["out1"]
	// sub2 的 instruction 包含 {out1}，应被替换
	cm1 := &stubChatModel{generateReply: model.ChatReply{Content: "first"}}
	cm2 := &stubChatModel{generateReply: model.ChatReply{Content: "got-first"}}

	sub1 := NewLLMAgent("a1", "inst1", "", "out1", cm1)
	sub2 := NewLLMAgent("a2", "instruction with {out1}", "", "", cm2)

	seq := NewSequentialAgent("seq", "desc", []model.Agent{sub1, sub2})
	result, err := seq.Run(context.Background(), testContent("hi"))

	assert.NoError(t, err)
	assert.Equal(t, "got-first", result)
}

func TestSequentialAgent_Run_SubAgentError_StopsExecution(t *testing.T) {
	cm1 := &stubChatModel{generateErr: fmt.Errorf("fail")}
	cm2 := &stubChatModel{generateReply: model.ChatReply{Content: "second"}}

	sub1 := NewLLMAgent("a1", "inst1", "", "", cm1)
	sub2 := NewLLMAgent("a2", "inst2", "", "", cm2)

	seq := NewSequentialAgent("seq", "desc", []model.Agent{sub1, sub2})
	_, err := seq.Run(context.Background(), testContent("hi"))

	assert.Error(t, err)
}

func TestSequentialAgent_Stream_LastAgentStreams(t *testing.T) {
	cm1 := &stubChatModel{generateReply: model.ChatReply{Content: "first"}}
	cm2 := &stubChatModel{streamDelta: "stream"}

	sub1 := NewLLMAgent("a1", "inst1", "", "out1", cm1)
	sub2 := NewLLMAgent("a2", "inst2", "", "", cm2)

	seq := NewSequentialAgent("seq", "desc", []model.Agent{sub1, sub2})
	out := make(chan string, 10)
	err := seq.Stream(context.Background(), testContent("hi"), out)
	close(out)

	assert.NoError(t, err)
	var texts []string
	for s := range out {
		texts = append(texts, s)
	}
	assert.Equal(t, []string{"s", "t", "r", "e", "a", "m"}, texts)
}

// ============================================================
// ParallelAgent 测试
// ============================================================

func TestParallelAgent_Run_ConcatenatesResults(t *testing.T) {
	cm1 := &stubChatModel{generateReply: model.ChatReply{Content: "aaa"}}
	cm2 := &stubChatModel{generateReply: model.ChatReply{Content: "bbb"}}

	sub1 := NewLLMAgent("a1", "inst1", "", "", cm1)
	sub2 := NewLLMAgent("a2", "inst2", "", "", cm2)

	par := NewParallelAgent("par", "desc", []model.Agent{sub1, sub2})
	result, err := par.Run(context.Background(), testContent("hi"))

	assert.NoError(t, err)
	assert.Contains(t, result, "[a1] aaa")
	assert.Contains(t, result, "[a2] bbb")
}

func TestParallelAgent_Run_SubAgentError_ReturnsError(t *testing.T) {
	cm1 := &stubChatModel{generateReply: model.ChatReply{Content: "ok"}}
	cm2 := &stubChatModel{generateErr: fmt.Errorf("fail")}

	sub1 := NewLLMAgent("a1", "inst1", "", "", cm1)
	sub2 := NewLLMAgent("a2", "inst2", "", "", cm2)

	par := NewParallelAgent("par", "desc", []model.Agent{sub1, sub2})
	_, err := par.Run(context.Background(), testContent("hi"))

	assert.Error(t, err)
}

func TestParallelAgent_Run_ConcurrentExecution(t *testing.T) {
	// 验证并发执行：两个 agent 都被调用
	cm1 := &stubChatModel{generateReply: model.ChatReply{Content: "first"}}
	cm2 := &stubChatModel{generateReply: model.ChatReply{Content: "second"}}

	sub1 := NewLLMAgent("a1", "inst1", "", "", cm1)
	sub2 := NewLLMAgent("a2", "inst2", "", "", cm2)

	par := NewParallelAgent("par", "desc", []model.Agent{sub1, sub2})
	result, err := par.Run(context.Background(), testContent("hi"))

	assert.NoError(t, err)
	// 结果应包含两个 agent 的输出
	assert.True(t, strings.Contains(result, "first"))
	assert.True(t, strings.Contains(result, "second"))
}

func TestParallelAgent_Stream_OutputsConcurrently(t *testing.T) {
	// 使用单次输出的 stub 避免逐字符并发竞争
	cm1 := &singleShotChatModel{content: "result-a"}
	cm2 := &singleShotChatModel{content: "result-b"}

	sub1 := NewLLMAgent("a1", "inst1", "", "", cm1)
	sub2 := NewLLMAgent("a2", "inst2", "", "", cm2)

	par := NewParallelAgent("par", "desc", []model.Agent{sub1, sub2})
	out := make(chan string, 20)
	err := par.Stream(context.Background(), testContent("hi"), out)
	close(out)

	assert.NoError(t, err)
	var texts []string
	for s := range out {
		texts = append(texts, s)
	}
	full := strings.Join(texts, "")
	assert.Contains(t, full, "[a1]")
	assert.Contains(t, full, "[a2]")
	assert.Contains(t, full, "result-a")
	assert.Contains(t, full, "result-b")
}

// ============================================================
// LoopAgent 测试
// ============================================================

func TestLoopAgent_Run_RepeatsSubAgents(t *testing.T) {
	cm := &stubChatModel{generateReply: model.ChatReply{Content: "tick"}}
	sub := NewLLMAgent("a1", "inst1", "", "", cm)

	loop := NewLoopAgent("loop", "desc", []model.Agent{sub}, 3)
	result, err := loop.Run(context.Background(), testContent("hi"))

	assert.NoError(t, err)
	assert.Contains(t, result, "[a1] tick")
}

func TestLoopAgent_Run_DefaultMaxIterations(t *testing.T) {
	cm := &stubChatModel{generateReply: model.ChatReply{Content: "ok"}}
	sub := NewLLMAgent("a1", "inst1", "", "", cm)

	loop := NewLoopAgent("loop", "desc", []model.Agent{sub}, 0) // 0 → 默认 3
	assert.Equal(t, 3, loop.maxIterations)
}

func TestLoopAgent_Run_SubAgentError_StopsLoop(t *testing.T) {
	callCount := 0
	errModel := &stubChatModel{}
	errModel.generateErr = fmt.Errorf("fail on call")

	// 用一个计数 stub
	countModel := &countingChatModel{reply: model.ChatReply{Content: "ok"}, failAfter: 2}
	sub := NewLLMAgent("a1", "inst1", "", "", countModel)

	loop := NewLoopAgent("loop", "desc", []model.Agent{sub}, 5)
	_, err := loop.Run(context.Background(), testContent("hi"))

	assert.Error(t, err)
	_ = callCount
	_ = errModel
}

func TestLoopAgent_Stream_ExecutesAndStreams(t *testing.T) {
	cm := &stubChatModel{streamDelta: "loop"}
	sub := NewLLMAgent("a1", "inst1", "", "", cm)

	loop := NewLoopAgent("loop", "desc", []model.Agent{sub}, 2)
	out := make(chan string, 20)
	err := loop.Stream(context.Background(), testContent("hi"), out)
	close(out)

	assert.NoError(t, err)
	var texts []string
	for s := range out {
		texts = append(texts, s)
	}
	full := strings.Join(texts, "")
	assert.Contains(t, full, "[a1]")
}

// ============================================================
// 工具函数测试
// ============================================================

func TestCloneVars_CreatesIndependentCopy(t *testing.T) {
	orig := map[string]string{"a": "1", "b": "2"}
	cloned := cloneVars(orig)
	cloned["c"] = "3"
	assert.NotContains(t, orig, "c")
}

func TestApplyVars_ReplacesPlaceholders(t *testing.T) {
	vars := map[string]string{"name": "world", "greeting": "hello"}
	result := applyVars("{greeting} {name}!", vars)
	assert.Equal(t, "hello world!", result)
}

func TestApplyVars_EmptyTemplate_ReturnsEmpty(t *testing.T) {
	result := applyVars("", map[string]string{"a": "1"})
	assert.Equal(t, "", result)
}

func TestApplyVars_NoVars_ReturnsOriginal(t *testing.T) {
	result := applyVars("hello {name}", nil)
	assert.Equal(t, "hello {name}", result)
}

func TestInitialMessages_WithInstruction(t *testing.T) {
	msgs := initialMessages("system instruction", "user text")
	assert.Len(t, msgs, 2)
	assert.Equal(t, model.ChatRoleSystem, msgs[0].Role)
	assert.Equal(t, "system instruction", msgs[0].Content)
	assert.Equal(t, model.ChatRoleUser, msgs[1].Role)
}

func TestInitialMessages_EmptyInstruction_SkipsSystem(t *testing.T) {
	msgs := initialMessages("", "user text")
	assert.Len(t, msgs, 1)
	assert.Equal(t, model.ChatRoleUser, msgs[0].Role)
}

func TestFirstText_WithContent(t *testing.T) {
	content := model.ChatContent{Texts: []model.TextPart{{Message: "hi"}, {Message: "bye"}}}
	assert.Equal(t, "hi", firstText(content))
}

func TestFirstText_EmptyContent(t *testing.T) {
	assert.Equal(t, "", firstText(model.ChatContent{}))
}

// ============================================================
// 辅助 stub
// ============================================================

// singleShotChatModel 一次性输出完整内容的 stub（适合并发测试）
type singleShotChatModel struct {
	content string
}

func (m *singleShotChatModel) Generate(ctx context.Context, msgs []model.ChatMessage) (model.ChatReply, error) {
	return model.ChatReply{Content: m.content}, nil
}

func (m *singleShotChatModel) Stream(ctx context.Context, msgs []model.ChatMessage) (<-chan model.ChatStreamEvent, <-chan error) {
	events := make(chan model.ChatStreamEvent, 2)
	errs := make(chan error, 1)
	go func() {
		events <- model.ChatStreamEvent{Delta: m.content, Done: true}
		close(events)
		close(errs)
	}()
	return events, errs
}

func (m *singleShotChatModel) CallTool(ctx context.Context, name, arguments string) (string, error) {
	return "ok", nil
}

// countingChatModel 记录调用次数，超过 failAfter 后返回错误
type countingChatModel struct {
	reply     model.ChatReply
	failAfter int
	calls     int
}

func (m *countingChatModel) Generate(ctx context.Context, msgs []model.ChatMessage) (model.ChatReply, error) {
	m.calls++
	if m.calls > m.failAfter {
		return model.ChatReply{}, fmt.Errorf("fail at call %d", m.calls)
	}
	return m.reply, nil
}

func (m *countingChatModel) Stream(ctx context.Context, msgs []model.ChatMessage) (<-chan model.ChatStreamEvent, <-chan error) {
	events := make(chan model.ChatStreamEvent, 4)
	errs := make(chan error, 1)
	m.calls++
	if m.calls > m.failAfter {
		errs <- fmt.Errorf("fail at call %d", m.calls)
	} else {
		events <- model.ChatStreamEvent{Delta: m.reply.Content, Done: true}
	}
	close(events)
	close(errs)
	return events, errs
}

func (m *countingChatModel) CallTool(ctx context.Context, name, arguments string) (string, error) {
	return "ok", nil
}