package agent

import (
	"bufio"
	"encoding/json"
	"errors"
	"io"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestCodexTranslatedOutputBackpressuresWithoutDroppingTranscriptOrTerminal(t *testing.T) {
	output := newCodexTranslatedOutput(1, 1024, 1024)
	defer output.close()

	transcript := `{"kind":"text","text":"first"}`
	terminal := `{"kind":"turn_end"}`
	require.NoError(t, output.enqueue(transcript, len(transcript)+1))

	terminalDone := make(chan error, 1)
	go func() {
		terminalDone <- output.enqueue(terminal, len(terminal)+1)
	}()
	select {
	case err := <-terminalDone:
		t.Fatalf("terminal enqueue did not back-pressure on a no-consumer queue: %v", err)
	case <-time.After(20 * time.Millisecond):
	}

	require.Equal(t, transcript, <-output.lines)
	output.release(len(transcript) + 1)
	require.NoError(t, <-terminalDone)
	require.Equal(t, terminal, <-output.lines)
	output.release(len(terminal) + 1)
	require.NoError(t, output.fatalError())
}

func TestCodexTranslatedOutputByteBudgetUsesFinalEncodedWireBytes(t *testing.T) {
	event := translatedEvent{
		Kind: consts.ParsedEventKindText,
		Text: `quotes " and slashes \\ occupy encoded bytes`,
	}
	line, wireBytes, err := encodeCodexTranslatedEvent(event)
	require.NoError(t, err)
	require.Equal(t, len(line)+1, wireBytes)
	require.Greater(t, wireBytes, len(event.Text), "budget must include JSON structure and escaping")

	output := newCodexTranslatedOutput(4, 1024, 1024)
	defer output.close()
	require.NoError(t, output.enqueue(line, wireBytes))
	output.mu.Lock()
	require.Equal(t, wireBytes, output.queuedBytes)
	output.mu.Unlock()
	require.Equal(t, line, <-output.lines)
	output.release(wireBytes)
}

func TestCodexTranslatedOutputDisablesHTMLExpansionBeforeParserLimitCheck(t *testing.T) {
	// json.Marshal expands each '<' from one byte to six. Choose a payload that
	// the old encoder pushed over Process's 32 MiB line limit while the actual
	// JSON line with HTML escaping disabled remains comfortably below it.
	text := strings.Repeat("<", codexTranslatedOutputMaxLineBytes/6+1024)
	event := translatedEvent{Kind: consts.ParsedEventKindText, Text: text}

	oldEncoding, err := json.Marshal(event)
	require.NoError(t, err)
	require.Greater(t, len(oldEncoding)+1, codexTranslatedOutputMaxLineBytes)

	line, wireBytes, err := encodeCodexTranslatedEvent(event)
	require.NoError(t, err)
	require.LessOrEqual(t, wireBytes, codexTranslatedOutputMaxLineBytes)
	require.Contains(t, line, "<<<")
	require.NotContains(t, line, `\u003c`)
}

func TestCodexTranslatedOutputOversizeIsTransportFatal(t *testing.T) {
	output := newCodexTranslatedOutput(1, 32, 64)
	fatal := make(chan error, 1)
	output.setFatalHandler(func(err error) { fatal <- err })

	err := output.enqueue(strings.Repeat("x", 32), 33)
	require.Error(t, err)
	require.Contains(t, err.Error(), "process parser limit")
	require.EqualError(t, <-fatal, err.Error())
	select {
	case <-output.done:
	default:
		t.Fatal("fatal overflow did not abort blocked generation producers")
	}
	output.close()
}

func TestCodexTransportResetAbortsFullOldOutputBeforeGenerationFence(t *testing.T) {
	oldOutput := newCodexTranslatedOutput(1, 1024, 1024)
	oldStdin := &codexTransportCloseRecorder{}
	driver := &CodexAppServerDriver{
		launchGeneration: 1,
		codexRPCState: codexRPCState{
			pending:          make(map[int]pendingAppServerRequest),
			serverRequests:   make(map[string]inFlightServerRequest),
			rawStdin:         oldStdin,
			translatedLines:  oldOutput.lines,
			translatedOutput: oldOutput,
		},
	}
	require.NoError(t, oldOutput.enqueue("first", len("first")+1))

	producerEntered := make(chan struct{})
	producerDone := make(chan error, 1)
	go func() {
		driver.transportGenerationMu.RLock()
		close(producerEntered)
		err := driver.enqueueTranslatedEvent(translatedEvent{
			Kind: consts.ParsedEventKindText,
			Text: "blocked-old-generation",
		})
		driver.transportGenerationMu.RUnlock()
		producerDone <- err
	}()
	<-producerEntered
	select {
	case err := <-producerDone:
		t.Fatalf("old generation producer did not block on the full queue: %v", err)
	case <-time.After(20 * time.Millisecond):
	}

	resetDone := make(chan struct{})
	var newOutput *codexTranslatedOutput
	go func() {
		_, _, newOutput = driver.installTransportGeneration(drivercontract.LaunchConfig{}, &codexTransportCloseRecorder{})
		close(resetDone)
	}()

	select {
	case err := <-producerDone:
		require.ErrorIs(t, err, errCodexTranslatedOutputSuperseded)
	case <-time.After(time.Second):
		t.Fatal("reset did not abort the old blocked output producer before taking the generation fence")
	}
	select {
	case <-resetDone:
	case <-time.After(time.Second):
		t.Fatal("generation reset remained blocked after aborting old output")
	}
	require.Equal(t, 1, oldStdin.closeCount())
	require.NotSame(t, oldOutput, newOutput)
	newOutput.close()
}

func TestCodexOldGenerationStderrCannotPolluteReplacementOutput(t *testing.T) {
	driver := &CodexAppServerDriver{}
	oldGeneration, _, _ := driver.installTransportGeneration(drivercontract.LaunchConfig{}, nil)
	newGeneration, newLines, newOutput := driver.installTransportGeneration(drivercontract.LaunchConfig{}, nil)
	require.Greater(t, newGeneration, oldGeneration)
	defer newOutput.close()

	driver.handleCodexAppServerStderrLineForGeneration(
		oldGeneration,
		"ERROR codex_app_server: Invalid configuration: Model provider stale not found",
	)
	require.Empty(t, newLines, "old stderr callback must not emit into replacement output")

	driver.handleCodexAppServerStderrLineForGeneration(
		newGeneration,
		"ERROR codex_app_server: Invalid configuration: Model provider current not found",
	)
	select {
	case line := <-newLines:
		events := driver.ParseLine(line)
		require.Len(t, events, 1)
		require.Equal(t, consts.ParsedEventKindError, events[0].Kind)
	case <-time.After(time.Second):
		t.Fatal("current generation stderr was not translated")
	}
}

func TestCodexTranslatedOutputAbortErrorIsNotFatal(t *testing.T) {
	output := newCodexTranslatedOutput(1, 1024, 1024)
	output.abort()
	err := output.enqueue("late", len("late")+1)
	require.ErrorIs(t, err, errCodexTranslatedOutputSuperseded)
	require.True(t, errors.Is(err, errCodexTranslatedOutputSuperseded))
	require.NoError(t, output.fatalError())
	output.close()
}

func TestCodexTranslatedOutputCancellationUnblocksNoConsumerPipeFlusher(t *testing.T) {
	driver := &CodexAppServerDriver{}
	output := newCodexTranslatedOutput(2, 1024, 1024)
	reader, writer := io.Pipe()
	defer reader.Close()
	output.attachWriter(writer)

	flusherDone := make(chan struct{})
	go func() {
		driver.flushTranslatedOutput(output, writer)
		close(flusherDone)
	}()
	require.NoError(t, output.enqueue("blocked-without-reader", len("blocked-without-reader")+1))
	require.Eventually(t, func() bool {
		return len(output.lines) == 0
	}, time.Second, time.Millisecond, "flusher never entered the no-consumer pipe write")

	output.cancel()
	select {
	case <-flusherDone:
	case <-time.After(time.Second):
		t.Fatal("generation cancellation did not close the translated io.PipeWriter and release its blocked flusher")
	}
	require.NoError(t, output.fatalError(), "intentional generation cancellation is not transport-fatal")
	output.close()
}

func TestCodexTranslatedOutputConcurrentAbortCloseAndFatalHandlerAreIdempotent(t *testing.T) {
	output := newCodexTranslatedOutput(2, 1024, 1024)
	reader, writer := io.Pipe()
	defer reader.Close()
	output.attachWriter(writer)
	output.setFatalHandler(func(error) { output.close() })

	var wg sync.WaitGroup
	for range 8 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			output.abort()
			output.close()
		}()
	}
	wg.Add(1)
	go func() {
		defer wg.Done()
		output.fail(errors.New("forced output failure"))
	}()
	wg.Wait()
	_, ok := <-output.lines
	require.False(t, ok)
}

func TestCodexTransportFailureClosesLegacyTranslatedLinesAfterDirectReplacement(t *testing.T) {
	managedOutput := newCodexTranslatedOutput(1, 1024, 1024)
	legacyLines := make(chan string, 1)
	aborted := make(chan struct{}, 1)
	driver := &CodexAppServerDriver{
		launchGeneration: 9,
		codexRPCState: codexRPCState{
			pending:          make(map[int]pendingAppServerRequest),
			serverRequests:   make(map[string]inFlightServerRequest),
			translatedOutput: managedOutput,
			// A number of focused legacy tests directly replace this field. The
			// managed queue must not make their replacement channel immortal.
			translatedLines: legacyLines,
			transportFailure: &codexAppServerTransportFailure{abort: func() {
				aborted <- struct{}{}
			}},
		},
	}
	failure := driver.transportFailure
	driver.failTransport(9, failure, "forced")

	select {
	case <-aborted:
	default:
		t.Fatal("transport failure did not invoke abort hook")
	}
	_, managedOpen := <-managedOutput.lines
	_, legacyOpen := <-legacyLines
	require.False(t, managedOpen)
	require.False(t, legacyOpen)
}

func TestCodexTransportEOFDrainsAdmittedTranscriptAndTerminalBeforeWriterClose(t *testing.T) {
	output := newCodexTranslatedOutput(4, 1024, 1024)
	reader, writer := io.Pipe()
	defer reader.Close()
	output.attachWriter(writer)
	aborted := make(chan struct{}, 1)
	failure := &codexAppServerTransportFailure{abort: func() { aborted <- struct{}{} }}
	driver := &CodexAppServerDriver{
		launchGeneration: 3,
		codexRPCState: codexRPCState{
			pending:          make(map[int]pendingAppServerRequest),
			serverRequests:   make(map[string]inFlightServerRequest),
			translatedLines:  output.lines,
			translatedOutput: output,
			transportFailure: failure,
		},
	}

	require.NoError(t, driver.enqueueTranslatedEvent(translatedEvent{
		Kind: consts.ParsedEventKindText,
		Text: "final transcript",
	}))
	require.NoError(t, driver.enqueueTranslatedEvent(translatedEvent{
		Kind: consts.ParsedEventKindTurnEnd,
	}))

	// Source EOF is transport-fatal, but it follows both frames on the ordered
	// app-server stream. Those already-admitted translated frames must drain
	// before the Process-facing writer closes.
	driver.readLoopForGeneration(strings.NewReader(""), 3, failure)

	// Start the flusher only after EOF sealed the queue. This deterministically
	// models a lagging Process consumer and proves that producer abort is not
	// mistaken for cancellation after the first successfully written line.
	flusherDone := make(chan struct{})
	go func() {
		driver.flushTranslatedOutput(output, writer)
		close(flusherDone)
	}()
	readDone := make(chan []string, 1)
	go func() {
		lineReader := bufio.NewReader(reader)
		var lines []string
		for {
			line, err := lineReader.ReadString('\n')
			if line != "" {
				lines = append(lines, strings.TrimSpace(line))
			}
			if err != nil {
				readDone <- lines
				return
			}
		}
	}()

	select {
	case lines := <-readDone:
		require.Len(t, lines, 2)
		first := driver.ParseLine(lines[0])
		second := driver.ParseLine(lines[1])
		require.Len(t, first, 1)
		require.Len(t, second, 1)
		require.Equal(t, consts.ParsedEventKindText, first[0].Kind)
		require.Equal(t, consts.ParsedEventKindTurnEnd, second[0].Kind)
	case <-time.After(time.Second):
		t.Fatal("transport EOF closed or stranded translated output before admitted terminal drain")
	}
	select {
	case <-flusherDone:
	case <-time.After(time.Second):
		t.Fatal("translated output flusher did not exit after sealed queue drain")
	}
	select {
	case <-aborted:
	default:
		t.Fatal("transport EOF did not terminate the bound app-server")
	}
}

func TestCodexTransportWaitCannotOvertakeReaderAndTranslatedTerminalDrain(t *testing.T) {
	output := newCodexTranslatedOutput(1, 1024, 1024)
	readLoopDone := make(chan struct{})
	translatedConsumed := make(chan struct{})
	childWaitCalled := make(chan struct{})
	waitDone := make(chan error, 1)
	go func() {
		waitDone <- waitForCodexAppServerTransport(
			&codexAppServerTransportFailure{},
			readLoopDone,
			output,
			translatedConsumed,
			func() error {
				close(childWaitCalled)
				return nil
			},
		)
	}()

	select {
	case <-childWaitCalled:
		t.Fatal("child wait overtook the stdout reader")
	case <-time.After(20 * time.Millisecond):
	}
	close(readLoopDone)
	select {
	case <-childWaitCalled:
		t.Fatal("child wait overtook translated terminal output")
	case <-time.After(20 * time.Millisecond):
	}
	output.finishFlush()
	select {
	case <-childWaitCalled:
		t.Fatal("child wait overtook Process terminal consumption")
	case <-time.After(20 * time.Millisecond):
	}
	close(translatedConsumed)
	select {
	case <-childWaitCalled:
	case <-time.After(time.Second):
		t.Fatal("child wait did not proceed after reader and translated output completed")
	}
	require.NoError(t, <-waitDone)
}

func TestCodexTransportWaitShutdownUnblocksMissingTranslatedConsumer(t *testing.T) {
	output := newCodexTranslatedOutput(1, 1024, 1024)
	output.finishFlush()
	readLoopDone := make(chan struct{})
	close(readLoopDone)
	translatedConsumed := make(chan struct{})
	failure := &codexAppServerTransportFailure{}
	waitDone := make(chan error, 1)
	go func() {
		waitDone <- waitForCodexAppServerTransport(
			failure,
			readLoopDone,
			output,
			translatedConsumed,
			func() error { return nil },
		)
	}()

	select {
	case <-waitDone:
		t.Fatal("transport wait ignored missing Process consumer before shutdown")
	case <-time.After(20 * time.Millisecond):
	}
	failure.beginShutdown()
	select {
	case err := <-waitDone:
		require.NoError(t, err)
	case <-time.After(time.Second):
		t.Fatal("explicit shutdown did not release transport wait with no Process consumer")
	}
}