package agent

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

	"ollie/backend"
)

// TestTruncation_LargeResultTruncated verifies that a tool result exceeding
// the 128KB safety limit is truncated and a hint is appended.
func TestTruncation_LargeResultTruncated(t *testing.T) {
	large := strings.Repeat("x", defaultToolResultMaxBytes+10_000)

	be := defaultBE()
	be.respond = toolsStream([]backend.ToolCall{
		{ID: "1", Name: "execute_code", Arguments: json.RawMessage(`{}`)},
	})

	c := newCore(t, be, nil)
	c.r.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) {
		return large, nil, nil
	}

	evs := collectEvents(context.Background(), c, "big result")

	toolEvs := byRole(evs, "tool")
	if len(toolEvs) != 1 {
		t.Fatalf("tool events = %d; want 1", len(toolEvs))
	}
	result := toolEvs[0]
	if !strings.HasPrefix(result, strings.Repeat("x", defaultToolResultMaxBytes)) {
		t.Errorf("result does not start with %d x's", defaultToolResultMaxBytes)
	}
	if !strings.Contains(result, "HARD LIMIT") {
		t.Errorf("truncation hint missing from result: %q", result[:min(len(result), 80)])
	}
	if !strings.Contains(result, "execute_code") {
		t.Errorf("tool name missing from truncation hint: %q", result[:min(len(result), 120)])
	}
}

// TestTruncation_SmallResultNotTruncated verifies that results within the limit
// pass through unchanged.
func TestTruncation_SmallResultNotTruncated(t *testing.T) {
	be := defaultBE()
	be.respond = toolsStream([]backend.ToolCall{
		{ID: "1", Name: "execute_code", Arguments: json.RawMessage(`{}`)},
	})

	c := newCore(t, be, nil)
	c.r.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) {
		return "short result", nil, nil
	}

	evs := collectEvents(context.Background(), c, "small result")

	toolEvs := byRole(evs, "tool")
	if len(toolEvs) != 1 {
		t.Fatalf("tool events = %d; want 1", len(toolEvs))
	}
	if toolEvs[0] != "short result" {
		t.Errorf("result = %q; want %q", toolEvs[0], "short result")
	}
}

// TestTruncation_ErrorTruncated verifies that error results are subject to the
// same 128KB safety limit (the original bug: errors bypassed all truncation).
func TestTruncation_ErrorTruncated(t *testing.T) {
	longErr := strings.Repeat("e", defaultToolResultMaxBytes+10_000)

	be := defaultBE()
	be.respond = toolsStream([]backend.ToolCall{
		{ID: "1", Name: "execute_code", Arguments: json.RawMessage(`{}`)},
	})

	c := newCore(t, be, nil)
	c.r.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) {
		return "", nil, &mockErr{longErr}
	}

	evs := collectEvents(context.Background(), c, "error result")

	toolEvs := byRole(evs, "tool")
	if len(toolEvs) != 1 {
		t.Fatalf("tool events = %d; want 1", len(toolEvs))
	}
	if !strings.Contains(toolEvs[0], "HARD LIMIT") {
		t.Errorf("error result was NOT truncated: len=%d", len(toolEvs[0]))
	}
	if len(toolEvs[0]) > defaultToolResultMaxBytes+200 {
		t.Errorf("error result too large after truncation: len=%d", len(toolEvs[0]))
	}
}

// TestTruncation_CachedResultAlreadyTruncated verifies that a cache hit on a
// previously-truncated result returns the truncated form, not the original.
func TestTruncation_CachedResultAlreadyTruncated(t *testing.T) {
	large := strings.Repeat("z", defaultToolResultMaxBytes+10_000)

	be := defaultBE()
	be.respond = multiTurnToolsStream([][]backend.ToolCall{
		{{ID: "1", Name: "file_read", Arguments: json.RawMessage(`{"path":"a.txt"}`)}},
		{{ID: "2", Name: "file_read", Arguments: json.RawMessage(`{"path":"a.txt"}`)}},
	})

	c := newCore(t, be, nil)
	c.r.runtime.ClassifyTool = func(string) bool { return true } // read-safe → cacheable
	c.r.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) {
		return large, nil, nil
	}

	evs := collectEvents(context.Background(), c, "cached truncated")

	toolEvs := byRole(evs, "tool")
	if len(toolEvs) != 2 {
		t.Fatalf("tool events = %d; want 2", len(toolEvs))
	}
	for i, ev := range toolEvs {
		if !strings.Contains(ev, "HARD LIMIT") {
			t.Errorf("event[%d] missing truncation hint: %q", i, ev[:min(len(ev), 80)])
		}
	}
}