package native

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/exec"
	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/llm"
)

// captureStream 捕获每次调用的 model / 首条 system 消息，并按序回放回合。
func captureStream(turns ...llm.Message) (llm.StreamFunc, *capture) {
	c := &capture{}
	var mu sync.Mutex
	i := 0
	return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 4)
		go func() {
			defer close(ch)
			mu.Lock()
			defer mu.Unlock()
			c.models = append(c.models, model)
			c.toolDecls = append(c.toolDecls, toolNames(tools))
			if len(messages) > 0 && messages[0].Role == llm.RoleSystem {
				c.system = messages[0].Content
			}
			if i >= len(turns) {
				ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}}
				return
			}
			msg := turns[i]
			i++
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg}
		}()
		return ch
	}, c
}

type capture struct {
	mu        sync.Mutex
	models    []string
	toolDecls [][]string
	system    string
}

func toolNames(decls []llm.ToolDecl) []string {
	var out []string
	for _, d := range decls {
		out = append(out, d.Name)
	}
	return out
}

func TestBackendRequestModelOverride(t *testing.T) {
	stream, c := captureStream()
	b := New(Options{
		ModelOptions: ModelOptions{
			Model:  "default-model",
			Stream: stream,
		},
	})
	dir := t.TempDir()
	res, err := b.Execute(context.Background(), exec.Request{
		WorkDir: dir, Prompt: "hi", Targets: []string{"output.md"}, Model: "override-model",
		ToolProfile: exec.ToolProfileReasoning,
	})
	if err != nil || !res.Success {
		t.Fatalf("execute: err=%v res=%+v", err, res)
	}
	if len(c.models) == 0 || c.models[0] != "override-model" {
		t.Fatalf("want override-model, got %v", c.models)
	}
}

func TestBackendRequestToolAllowlist(t *testing.T) {
	stream, c := captureStream(
		llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
			ID: "c1", Name: "write_file", Arguments: `{"path":"output.md","content":"blocked"}`,
		}}},
		llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
			ID: "c2", Name: "write_file", Arguments: `{"path":"output2.md","content":"still blocked"}`,
		}}},
		llm.Message{Role: llm.RoleAssistant, Content: "gave up", StopReason: llm.StopEnd},
	)
	b := New(Options{
		ModelOptions: ModelOptions{
			Model:  "fake",
			Stream: stream,
		},
	})
	dir := t.TempDir()
	// 目标 output2.md 不在 allowlist 内能写吗？——write_file 被整体收窄，
	// 模型两次尝试都会得到 unknown tool，最终无交付物 → 失败。
	res, err := b.Execute(context.Background(), exec.Request{
		WorkDir: dir, Prompt: "写文件", Targets: []string{"output2.md"},
		ToolAllowlist: []string{"read_file", "list_dir", "grep"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if res.Success {
		t.Fatalf("write_file restricted out; deliverable must fail, got %+v", res)
	}
	if len(c.toolDecls) == 0 {
		t.Fatal("expected tool decls")
	}
	for _, names := range c.toolDecls {
		for _, n := range names {
			if n == "write_file" || n == "edit_file" {
				t.Fatalf("write/edit must be restricted out of decls, got %v", names)
			}
		}
	}
}

func TestBackendRequestUnknownAllowlistFailsClosed(t *testing.T) {
	b := New(Options{
		ModelOptions: ModelOptions{
			Model:  "fake",
			Stream: script(),
		},
	})
	res, err := b.Execute(context.Background(), exec.Request{
		WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
		ToolAllowlist: []string{"read_file", "no_such_tool"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if res.Success || !strings.Contains(res.ErrorReason, "no_such_tool") {
		t.Fatalf("want fail-closed unknown allowlist name, got %+v", res)
	}
}

func TestBackendRequestBudgetDeadlineBindsContext(t *testing.T) {
	deadlineHit := make(chan struct{})
	var once sync.Once
	stream := func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 4)
		go func() {
			defer close(ch)
			select {
			case <-ctx.Done():
				once.Do(func() { close(deadlineHit) })
				ch <- llm.Event{Kind: llm.KindError, Err: ctx.Err(), Message: llm.Message{StopReason: llm.StopAborted}}
			case <-time.After(3 * time.Second):
				ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}}
			}
		}()
		return ch
	}
	b := New(Options{
		ModelOptions: ModelOptions{
			Model:  "fake",
			Stream: stream,
		},
	})
	start := time.Now()
	res, err := b.Execute(context.Background(), exec.Request{
		WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
		Timeout: 30 * time.Second, // 长超时，预算截止应更先生效
		Budget:  &exec.Budget{Deadline: time.Now().Add(150 * time.Millisecond)},
	})
	if err != nil {
		t.Fatal(err)
	}
	select {
	case <-deadlineHit:
	case <-time.After(2 * time.Second):
		t.Fatal("budget deadline should bind the run context")
	}
	if res.Success {
		t.Fatalf("deadline run must not succeed, got %+v", res)
	}
	if elapsed := time.Since(start); elapsed > 5*time.Second {
		t.Fatalf("deadline should fire early, took %s", elapsed)
	}
}

func TestBackendRequestExecutionIDInLogs(t *testing.T) {
	stream, _ := captureStream()
	b := New(Options{
		ModelOptions: ModelOptions{
			Model:  "fake",
			Stream: stream,
		},
	})
	dir := t.TempDir()
	res, err := b.Execute(context.Background(), exec.Request{
		WorkDir: dir, Prompt: "hi", Targets: []string{"output.md"},
		Instructions: "你是角色",
		ExecutionID:  "exec-42", SessionScope: "run-9/step-3",
		ToolProfile: exec.ToolProfileReasoning,
	})
	if err != nil || !res.Success {
		t.Fatalf("execute: err=%v res=%+v", err, res)
	}
	logBytes, _ := os.ReadFile(filepath.Join(dir, "execution.log"))
	logs := string(logBytes)
	if !strings.Contains(logs, "execution_id=exec-42") || !strings.Contains(logs, "run-9/step-3") {
		t.Fatalf("execution log should carry ids, got:\n%s", logs)
	}
}