package agent

import (
	"bufio"
	"bytes"
	"context"
	"io"
	"log/slog"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
)

func TestReadLineWithLimitReturnsCompleteLineWithNewline(t *testing.T) {
	reader := bufio.NewReader(strings.NewReader("hello\nrest"))

	line, err := readLineWithLimit(reader, 16)

	require.NoError(t, err)
	require.Equal(t, "hello\n", line)
}

func TestReadLineWithLimitReturnsPartialLineAtEOF(t *testing.T) {
	reader := bufio.NewReader(strings.NewReader("unterminated"))

	line, err := readLineWithLimit(reader, 16)

	require.NoError(t, err)
	require.Equal(t, "unterminated", line)
}

func TestReadLineWithLimitReturnsEOFWhenNoBytesRemain(t *testing.T) {
	reader := bufio.NewReader(strings.NewReader(""))

	line, err := readLineWithLimit(reader, 16)

	require.ErrorIs(t, err, io.EOF)
	require.Empty(t, line)
}

func TestReadLineWithLimitRejectsLineOverLimit(t *testing.T) {
	reader := bufio.NewReader(strings.NewReader("12345\n"))

	line, err := readLineWithLimit(reader, 4)

	require.Error(t, err)
	require.Contains(t, err.Error(), "message exceeds 4 bytes")
	require.Empty(t, line)
}

func TestReadLineWithLimitStopsReadingAnUnterminatedOversizedStream(t *testing.T) {
	source := &countingFillReader{remaining: 1 << 20}
	reader := bufio.NewReaderSize(source, 16)

	line, err := readLineWithLimit(reader, 32)

	require.ErrorContains(t, err, "message exceeds 32 bytes")
	require.Empty(t, line)
	require.LessOrEqual(t, source.read, 48,
		"the line limit must be enforced while accumulating, not after buffering the complete hostile stream")
}

type countingFillReader struct {
	remaining int
	read      int
}

func (r *countingFillReader) Read(p []byte) (int, error) {
	if r.remaining <= 0 {
		return 0, io.EOF
	}
	n := len(p)
	if n > r.remaining {
		n = r.remaining
	}
	for i := 0; i < n; i++ {
		p[i] = 'x'
	}
	r.remaining -= n
	r.read += n
	return n, nil
}

func TestParseStdoutDoesNotLogRawRuntimeLine(t *testing.T) {
	var logs bytes.Buffer
	previousLogger := slog.Default()
	slog.SetDefault(slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{
		Level: slog.LevelInfo,
	})))
	t.Cleanup(func() {
		slog.SetDefault(previousLogger)
	})

	ap := &Process{
		AgentID:               "agent-log-redaction",
		processRuntimeBinding: processRuntimeBinding{driver: &parseStdoutLogStubDriver{}}}

	ap.parseStdout(strings.NewReader(`{"kind":"text","text":"secret message content"}` + "\n"))

	require.NotContains(t, logs.String(), "secret message content")
	require.NotContains(t, logs.String(), "agent stdout")
}

type parseStdoutLogStubDriver struct{}

func (d *parseStdoutLogStubDriver) Name() string { return "test" }

func (d *parseStdoutLogStubDriver) SupportsStdinNotification() bool { return false }

func (d *parseStdoutLogStubDriver) MCPToolPrefix() string { return "" }

func (d *parseStdoutLogStubDriver) Spawn(context.Context, drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) {
	return drivercontract.SpawnResult{}, nil
}

func (d *parseStdoutLogStubDriver) ParseLine(string) []drivercontract.ParsedEvent { return nil }

func (d *parseStdoutLogStubDriver) EncodeStdinMessage(text, _ string) []byte {
	return []byte(text)
}