package main

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
	"path/filepath"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestNewTurnKeepsImagesInEveryChainRequestButNotLaterTurns(t *testing.T) {
	server := newModelFake(t)
	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		if call.number == 0 {
			modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "date", `{"id":0}`)}, 100, 0, 5))

			return
		}

		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("A red square.")}, 100, 0, 5))
	}

	test := newTestTurnLoop(t, server)
	data := imageTestPNG(t, 3, 3)
	name, err := test.agent.Images().Save(data)
	require.NoError(t, err)

	marker := imageText([]string{name}, "What is this?")
	_, err = test.loop.Submit(terminalSource, marker)
	require.NoError(t, err)
	test.events.WaitFor(t, "idle", 1)
	require.Equal(t, 2, server.Calls())

	first := server.Call(0).body["input"].([]any)
	second := server.Call(1).body["input"].([]any)
	require.Equal(t, first, second[:len(first)])

	expected := "data:image/png;base64," + base64.StdEncoding.EncodeToString(data)
	require.Contains(t, string(server.Call(0).raw), expected)
	require.Contains(t, string(server.Call(1).raw), expected)

	_, err = test.loop.Submit(terminalSource, "Continue with text.")
	require.NoError(t, err)
	test.events.WaitFor(t, "idle", 2)
	require.NotContains(t, string(server.Call(2).raw), "input_image")
	require.NotContains(t, strings.Join(test.Log(t), "\n"), "base64")
}

func TestImageArrivingDuringAChainJoinsItsNextRequest(t *testing.T) {
	server := newModelFake(t)
	arrived := make(chan struct{})
	release := make(chan struct{})
	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		if call.number == 0 {
			close(arrived)
			<-release
			modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "date", `{"id":0}`)}, 100, 0, 5))

			return
		}

		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("done")}, 100, 0, 5))
	}

	test := newTestTurnLoop(t, server)
	_, err := test.loop.Submit(terminalSource, "Start.")
	require.NoError(t, err)
	<-arrived

	name, err := test.agent.Images().Save(imageTestPNG(t, 2, 2))
	require.NoError(t, err)

	_, err = test.loop.Submit(terminalSource, imageText([]string{name}, "Read this."))
	require.NoError(t, err)
	close(release)
	test.events.WaitFor(t, "idle", 1)
	require.NotContains(t, string(server.Call(0).raw), "input_image")
	require.Contains(t, string(server.Call(1).raw), "input_image")
}

func TestLeafCompactionDescribesImagesOnceAndMergesOnlyText(t *testing.T) {
	server := newModelFake(t)
	store := &ImageStore{dir: filepath.Join(t.TempDir(), imageDirName)}
	name, err := store.Save(imageTestPNG(t, 2, 2))
	require.NoError(t, err)

	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		text := "receipt, Ada, $12, red square"
		if call.number == 0 {
			text = strings.Repeat("detail ", 100)
		}

		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(text)}, 100, 0, 5))
	}

	compactor := newTestCompactor(t, server)
	compactor.images = store
	task := Task{
		Level:   0,
		Index:   0,
		Context: "<chat>\n</chat>",
		Input:   []string{"user: " + imageText([]string{name}, "Read the receipt.")},
		Images:  []string{name},
	}
	text, err := compactor.Summarize(t.Context(), task)
	require.NoError(t, err)
	require.Equal(t, "receipt, Ada, $12, red square", text)
	require.Contains(t, string(server.Call(0).raw), "input_image")
	require.Contains(t, string(server.Call(0).raw), "kind, visible text, amounts, names, and context")
	require.NotContains(t, string(server.Call(1).raw), "input_image")

	_, err = compactor.Summarize(t.Context(), Task{
		Level:   1,
		Index:   0,
		Context: "<chat>\n</chat>",
		Input:   []string{text, "user: next"},
	})
	require.NoError(t, err)
	require.NotContains(t, string(server.Call(2).raw), "input_image")

	memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, compactor, discardLog())
	require.NoError(t, err)
	t.Cleanup(func() {
		require.NoError(t, memory.Close())
	})

	_, err = memory.Append("user", imageText([]string{name}, "A short caption."))
	require.NoError(t, err)
	require.NoError(t, memory.Settle(t.Context()))
	require.Equal(t, 4, server.Calls())
	require.Contains(t, string(server.Call(3).raw), "input_image")
	require.Contains(t, memory.View(), "receipt, Ada, $12")
}

func TestZoomReturnsImagePartsButLogsOnlyText(t *testing.T) {
	server := newModelFake(t)
	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		if call.number == 0 {
			modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "zoom", `{"id":0,"n":1,"name":null}`)}, 100, 0, 5))

			return
		}

		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("A red square.")}, 100, 0, 5))
	}

	test := newTestTurnLoop(t, server)
	name, err := test.agent.Images().Save(imageTestPNG(t, 2, 2))
	require.NoError(t, err)

	_, err = test.memory.Append("user", imageText([]string{name}, "Old image."))
	require.NoError(t, err)

	_, err = test.loop.Submit(terminalSource, "Zoom the old image.")
	require.NoError(t, err)
	test.events.WaitFor(t, "idle", 1)
	require.NotContains(t, string(server.Call(0).raw), "input_image")

	input := server.Call(1).body["input"].([]any)
	result := input[len(input)-1].(map[string]any)
	require.Equal(t, "function_call_output", result["type"])

	parts := result["output"].([]any)
	require.Len(t, parts, 2)
	require.Equal(t, "input_text", parts[0].(map[string]any)["type"])
	require.Equal(t, "input_image", parts[1].(map[string]any)["type"])
	require.Contains(t, parts[0].(map[string]any)["text"], "Old image.")
	require.NotContains(t, strings.Join(test.Log(t), "\n"), "base64")
}

func TestImageUserMessageKeepsTextPartsAndMarkerOrder(t *testing.T) {
	store := &ImageStore{dir: t.TempDir()}
	first, err := store.Save(imageTestPNG(t, 1, 2))
	require.NoError(t, err)

	second, err := store.Save(imageTestPNG(t, 2, 1))
	require.NoError(t, err)

	text := imageText([]string{first, second}, "Compare.")
	item, err := imageUserMessage(store, text)
	require.NoError(t, err)

	data, err := json.Marshal(item)
	require.NoError(t, err)

	var message struct {
		Role    string           `json:"role"`
		Content []ModelInputPart `json:"content"`
	}

	require.NoError(t, json.Unmarshal(data, &message))
	require.Equal(t, "user", message.Role)
	require.Len(t, message.Content, 3)
	require.Equal(t, text, message.Content[0].Text)
	require.NotEqual(t, message.Content[1].ImageURL, message.Content[2].ImageURL)
}