// A streamed, unfinished JSON tool argument still needs to expose code for live progress; E2E timing cannot reliably stop mid-argument.
package daemon

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

func TestExtractPartialJSONCode(t *testing.T) {
	code1 := ExtractPartialJSONCode(`{"code":"read('second.`)
	if code1 != "read('second." {
		t.Fatalf("expected read('second., got %s", code1)
	}

	code2 := ExtractPartialJSONCode(`{"code":"from pathlib import Path\nPath('demo.py').write_text('hello')"}`)
	expected2 := "from pathlib import Path\nPath('demo.py').write_text('hello')"
	if code2 != expected2 {
		t.Fatalf("expected %s, got %s", expected2, code2)
	}
}

func TestToolProgressReporterPreview(t *testing.T) {
	tests := []struct {
		want      *ToolCodePreview
		name      string
		arguments string
	}{
		{name: "empty", arguments: `{"code":""}`},
		{name: "short", arguments: `{"code":"print('hello')"}`, want: &ToolCodePreview{Text: "print('hello')"}},
		{name: "512 runes", arguments: `{"code":"` + strings.Repeat("a", 512) + `"}`, want: &ToolCodePreview{Text: strings.Repeat("a", 512)}},
		{name: "513 runes", arguments: `{"code":"x` + strings.Repeat("a", 512) + `"}`, want: &ToolCodePreview{Offset: 1, Text: strings.Repeat("a", 512)}},
		{name: "multibyte", arguments: `{"code":"é界🙂"}`, want: &ToolCodePreview{Text: "é界🙂"}},
		{name: "long unicode", arguments: `{"code":"` + strings.Repeat("é", 600) + strings.Repeat("界🙂", 256) + `"}`, want: &ToolCodePreview{Offset: 600, Text: strings.Repeat("界🙂", 256)}},
		{name: "control and formatting", arguments: `{"code":"a\n\t\r\b\f\u0000\u007f\u0085\u200b\u202eb"}`, want: &ToolCodePreview{Text: "a" + strings.Repeat(" ", 10) + "b"}},
		{name: "sanitized tail offset", arguments: `{"code":"\n` + strings.Repeat("a", 511) + `\u200b"}`, want: &ToolCodePreview{Offset: 1, Text: strings.Repeat("a", 511) + " "}},
		{name: "malformed utf8", arguments: "{\"code\":\"a\xff\xc0\xafz\"}", want: &ToolCodePreview{Text: "a���z"}},
		{name: "malformed utf8 tail", arguments: "{\"code\":\"\xff" + strings.Repeat("界", 512) + "\"}", want: &ToolCodePreview{Offset: 1, Text: strings.Repeat("界", 512)}},
		{name: "escaped json", arguments: `{"code":"print(\"é\")\\path\/\u754c\ud83d\ude42"}`, want: &ToolCodePreview{Text: "print(\"é\")\\path/界🙂"}},
		{name: "unfinished string", arguments: `{"code":"read('second.`, want: &ToolCodePreview{Text: "read('second."}},
		{name: "unfinished escape", arguments: `{"code":"hello\`, want: &ToolCodePreview{Text: "hello"}},
		{name: "unfinished unicode escape", arguments: `{"code":"hello\u75`, want: &ToolCodePreview{Text: "hello"}},
		{name: "unfinished utf8", arguments: "{\"code\":\"hello\xe7\x95", want: &ToolCodePreview{Text: "hello"}},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var published *ToolProgress
			reporter := NewToolProgressReporter(func(progress *ToolProgress) error {
				published = progress
				return nil
			})
			call := &ToolCallAssembly{ID: "preview"}
			call.Function.Name = "python"
			call.Function.Arguments = test.arguments
			if err := reporter.Report(call, "generating"); err != nil {
				t.Fatal(err)
			}
			if published == nil {
				t.Fatal("no progress published")
			}
			if test.want == nil {
				if published.Code != nil {
					t.Fatalf("expected no preview, got %+v", published.Code)
				}
				return
			}
			if published.Code == nil || *published.Code != *test.want {
				t.Fatalf("preview = %+v, want %+v", published.Code, test.want)
			}
		})
	}
}

var benchmarkPreviewBytes int

func BenchmarkToolProgressGrowingPythonArguments(b *testing.B) {
	for _, payload := range []struct {
		name string
		size int
	}{
		{name: "16KiB", size: 16 * 1024},
		{name: "128KiB", size: 128 * 1024},
	} {
		b.Run(payload.name, func(b *testing.B) {
			arguments, err := json.Marshal(map[string]string{"code": "#" + strings.Repeat("a", payload.size-1)})
			if err != nil {
				b.Fatal(err)
			}
			encoded := string(arguments)
			var prefixes []string
			for end := 1024; end < len(encoded); end += 1024 {
				prefixes = append(prefixes, encoded[:end])
			}
			prefixes = append(prefixes, encoded)
			publish := func(progress *ToolProgress) error {
				if progress.Code != nil {
					benchmarkPreviewBytes += len(progress.Code.Text)
				}
				return nil
			}
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				reporter := NewToolProgressReporter(publish)
				call := &ToolCallAssembly{ID: "benchmark"}
				call.Function.Name = "python"
				for _, prefix := range prefixes {
					call.Function.Arguments = prefix
					if err := reporter.Report(call, "generating"); err != nil {
						b.Fatal(err)
					}
				}
			}
		})
	}
}