package loop

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/guardrail"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

func TestRun_NoTools_FinalText(t *testing.T) {
	t.Parallel()
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Model: "m", Content: "Hello!"}}}}
	a := newTestAgent(t, ad, nil)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess-1", []model.Message{model.TextMessage(model.RoleUser, "hi")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "Hello!" {
		t.Fatalf("Text = %q, want Hello!", res.Text)
	}
	if res.Steps != 1 {
		t.Fatalf("Steps = %d, want 1", res.Steps)
	}
	if res.StopReason != agent.StopReasonStop {
		t.Fatalf("StopReason = %q, want %q", res.StopReason, agent.StopReasonStop)
	}
	if res.ToolCallsExecuted != 0 {
		t.Fatalf("ToolCallsExecuted = %d, want 0", res.ToolCallsExecuted)
	}
}

func TestRun_ToolCallThenFinal_FeedsBack(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"hi"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "Echoed!"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "echo this")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "Echoed!" {
		t.Fatalf("Text = %q, want Echoed!", res.Text)
	}
	if res.Steps != 2 {
		t.Fatalf("Steps = %d, want 2", res.Steps)
	}
	if res.StopReason != agent.StopReasonStop {
		t.Fatalf("StopReason = %q, want stop", res.StopReason)
	}
	if res.ToolCallsExecuted != 1 {
		t.Fatalf("ToolCallsExecuted = %d, want 1", res.ToolCallsExecuted)
	}
	if len(ad.reqs) < 2 {
		t.Fatalf("adapter requests = %d, want >= 2", len(ad.reqs))
	}
	last := ad.reqs[len(ad.reqs)-1]
	got := last.Messages[len(last.Messages)-1]
	if got.Role != model.RoleTool {
		t.Fatalf("last message role = %q, want tool", got.Role)
	}
	if !strings.Contains(got.Content[0].Text, "hi") {
		t.Fatalf("tool result text = %q, want to contain 'hi'", got.Content[0].Text)
	}
}

func TestRun_MaxRounds(t *testing.T) {
	t.Parallel()
	// 每轮使用不同的调用，使关键工具循环安全网不会
	// 介入；这里的目标是演练轮次预算。
	call1 := `{"tool_calls":[{"name":"echo","args":{"q":"x"}}]}`
	call2 := `{"tool_calls":[{"name":"echo","args":{"q":"y"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call1}},
		{resp: model.Response{Content: call2}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "x"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 2})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Steps != 2 {
		t.Fatalf("Steps = %d, want 2", res.Steps)
	}
	if res.StopReason != StopReasonMaxRounds {
		t.Fatalf("StopReason = %q, want %q", res.StopReason, StopReasonMaxRounds)
	}
	if res.ToolCallsExecuted != 2 {
		t.Fatalf("ToolCallsExecuted = %d, want 2", res.ToolCallsExecuted)
	}
}

func TestRun_ToolError_FeedsBack(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"boom","args":{}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "recovered"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "boom", err: errors.New("kaboom")})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "boom now")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "recovered" {
		t.Fatalf("Text = %q, want recovered", res.Text)
	}
	if res.ToolCallsExecuted != 1 {
		t.Fatalf("ToolCallsExecuted = %d, want 1", res.ToolCallsExecuted)
	}
	got := ad.reqs[len(ad.reqs)-1].Messages[len(ad.reqs[len(ad.reqs)-1].Messages)-1]
	if !strings.Contains(got.Content[0].Text, "kaboom") {
		t.Fatalf("tool result should carry the error, got %q", got.Content[0].Text)
	}
}

func TestRun_ToolTimeout_FeedsBack(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"slow","args":{}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "continued after timeout"}},
	}}
	td := dispatcherWithTimeout(t, 30*time.Millisecond, &mockTool{name: "slow", delay: 500 * time.Millisecond})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "slow it")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "continued after timeout" {
		t.Fatalf("Text = %q, want recovery text", res.Text)
	}
	got := ad.reqs[len(ad.reqs)-1].Messages[len(ad.reqs[len(ad.reqs)-1].Messages)-1]
	if !strings.Contains(got.Content[0].Text, "timed out") {
		t.Fatalf("tool result should report the timeout, got %q", got.Content[0].Text)
	}
}

func TestRun_CheckpointPerStep(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"hi"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	cp := newTestCheckpointStore(t)
	l := New(Options{Agent: a, Checkpoint: cp, MaxRounds: 8})

	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	got, err := cp.List(context.Background(), "sess")
	if err != nil {
		t.Fatalf("List: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("checkpoints = %d, want 2", len(got))
	}
	if got[0].Step != 1 || got[1].Step != 2 {
		t.Fatalf("checkpoint steps = [%d,%d], want [1,2]", got[0].Step, got[1].Step)
	}
}

func TestRun_TraceSpans(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"hi"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	var buf bytes.Buffer
	tr := trace.New(trace.NewJSONRecorder(&buf))
	l := New(Options{Agent: a, Tracer: tr, MaxRounds: 8})

	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	saw := map[trace.SpanKind]bool{}
	for _, r := range decodeSpanLines(t, buf.Bytes()) {
		saw[r.Kind] = true
	}
	for _, k := range []trace.SpanKind{trace.SpanKindLLM, trace.SpanKindTool, trace.SpanKindStateTransition} {
		if !saw[k] {
			t.Fatalf("missing span kind %q in %+v", k, saw)
		}
	}
}

func TestRun_MemoryInject(t *testing.T) {
	t.Parallel()
	store := newTestMemoryStore(t)
	if _, err := store.Add(context.Background(), memory.MemoryItem{
		PersonaID: "persona-a", Content: "她喜欢喝冰美式",
		Scope: memory.ScopePersona, Kind: memory.KindRelationship,
	}); err != nil {
		t.Fatalf("seed memory: %v", err)
	}
	ret := memory.NewRetriever(store)
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "喜欢冰美式"}}}}
	a := newTestAgent(t, ad, nil)
	l := New(Options{Agent: a, Retriever: ret, PersonaID: "persona-a", MaxRounds: 8})

	if _, err := l.Run(context.Background(), "sess", nil); err != nil {
		t.Fatalf("Run: %v", err)
	}
	// 空对话仍会走 injectMemory；查询为空时 retriever
	// 按最近程度返回该角色的条目。
	if len(ad.reqs) == 0 {
		t.Fatal("no request captured")
	}
	found := false
	for _, m := range ad.reqs[0].Messages {
		if m.Role == model.RoleSystem && strings.Contains(m.Content[0].Text, "冰美式") {
			found = true
			break
		}
	}
	if !found {
		t.Fatalf("injected memory not found in request messages: %+v", ad.reqs[0].Messages)
	}
}

func TestRun_MemoryExtract(t *testing.T) {
	t.Parallel()
	store := newTestMemoryStore(t)
	ext := memory.NewManualExtractor("persona-a", "role-1")
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "记住了"}}}}
	a := newTestAgent(t, ad, nil)
	l := New(Options{Agent: a, Memory: store, Extractor: ext, PersonaID: "persona-a", MaxRounds: 8})

	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "记住 我爱喝冰美式")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	items, err := store.ListByPersona(context.Background(), "persona-a")
	if err != nil {
		t.Fatalf("ListByPersona: %v", err)
	}
	if len(items) != 1 {
		t.Fatalf("items = %d, want 1", len(items))
	}
	if items[0].Content != "我爱喝冰美式" {
		t.Fatalf("item content = %q, want 我爱喝冰美式", items[0].Content)
	}
}

func TestRun_StreamDeltas(t *testing.T) {
	t.Parallel()
	ad := &mockAdapter{
		stream: []mgResp{{resp: model.Response{Content: "Hello world"}}},
		chunks: [][]string{{"Hel", "lo ", "world"}},
	}
	a := newTestAgent(t, ad, nil)
	var mu sync.Mutex
	var deltas []string
	l := New(Options{
		Agent: a,
		OnStream: func(ev model.StreamEvent) {
			if ev.Kind != model.StreamKindTextDelta {
				return
			}
			mu.Lock()
			deltas = append(deltas, ev.Delta)
			mu.Unlock()
		},
		MaxRounds: 8,
	})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "hi")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "Hello world" {
		t.Fatalf("Text = %q, want Hello world", res.Text)
	}
	mu.Lock()
	got := append([]string(nil), deltas...)
	mu.Unlock()
	if len(got) != 3 || got[0] != "Hel" || got[1] != "lo " || got[2] != "world" {
		t.Fatalf("deltas = %+v, want [Hel lo  world]", got)
	}
}

func TestRun_StreamGuardrailBlocks(t *testing.T) {
	t.Parallel()
	g := &mockGuard{inputVerdict: &guardrail.Verdict{Allow: false, Reason: "blocked in"}}
	pipe := guardrail.New(guardrail.WithGuard(g))
	ad := &mockAdapter{}
	a := newTestAgent(t, ad, nil)
	l := New(Options{Agent: a, Guards: pipe, OnStream: func(model.StreamEvent) {}, MaxRounds: 8})

	res, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "hi")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.StopReason != agent.StopReasonGuardrailInput {
		t.Fatalf("StopReason = %q, want %q", res.StopReason, agent.StopReasonGuardrailInput)
	}
	if res.Text != "blocked in" {
		t.Fatalf("Text = %q, want blocked in", res.Text)
	}
}

func TestRun_ContextCanceled(t *testing.T) {
	t.Parallel()
	ad := &mockAdapter{}
	a := newTestAgent(t, ad, nil)
	l := New(Options{Agent: a, MaxRounds: 8})

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := l.Run(ctx, "sess", []model.Message{model.TextMessage(model.RoleUser, "hi")}); !errors.Is(err, context.Canceled) {
		t.Fatalf("Run error = %v, want context.Canceled", err)
	}
}

func TestRun_NilAgent(t *testing.T) {
	t.Parallel()
	l := New(Options{})
	if _, err := l.Run(context.Background(), "sess", nil); !errors.Is(err, ErrNilAgent) {
		t.Fatalf("Run error = %v, want ErrNilAgent", err)
	}
}

func TestNew_DefaultMaxRounds(t *testing.T) {
	t.Parallel()
	if got := New(Options{}).MaxRounds(); got != DefaultMaxRounds {
		t.Fatalf("MaxRounds = %d, want %d", got, DefaultMaxRounds)
	}
	if got := New(Options{MaxRounds: 3}).MaxRounds(); got != 3 {
		t.Fatalf("MaxRounds = %d, want 3", got)
	}
}

// decodeSpanLines 把作为 JSONL 发出的每个 JSON span 读入 recs，
// 任何一行格式错误都让测试失败。
func decodeSpanLines(t *testing.T, data []byte) []trace.SpanRecord {
	t.Helper()
	sc := bufio.NewScanner(bytes.NewReader(data))
	var recs []trace.SpanRecord
	for sc.Scan() {
		line := sc.Bytes()
		if len(bytes.TrimSpace(line)) == 0 {
			continue
		}
		var r trace.SpanRecord
		if err := json.Unmarshal(line, &r); err != nil {
			t.Fatalf("decode span: %v", err)
		}
		recs = append(recs, r)
	}
	if err := sc.Err(); err != nil {
		t.Fatalf("scan spans: %v", err)
	}
	return recs
}