package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/pelletier/go-toml/v2"
	"github.com/stretchr/testify/require"

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

func TestKimiACPPrepareLaunchInjectsProviderWithoutLogin(t *testing.T) {
	stateDir := t.TempDir()
	prepared, err := prepareKimiACPLaunch(drivercontract.LaunchConfig{
		WorkDir:         t.TempDir(),
		RuntimeStateDir: stateDir,
		PermissionMode:  types.AgentPermissionModeFullAccess,
		Model:           "moonshot-v1",
		Provider: &drivercontract.ProviderConfig{
			Name: "AgentTeam Gateway", BaseURL: "https://llm.example.com/v1", APIKey: "sk-secret",
			WireAPIs: []string{provider.WireAPIAnthropic},
		},
	})
	require.NoError(t, err)
	require.Equal(t, []string{"acp"}, prepared.Args)
	env := strings.Join(prepared.Env, "\n")
	require.Contains(t, env, "KIMI_MODEL_NAME=moonshot-v1")
	require.Contains(t, env, "KIMI_MODEL_API_KEY=sk-secret")
	require.Contains(t, env, "KIMI_MODEL_BASE_URL=https://llm.example.com/v1")
	require.Contains(t, env, "KIMI_MODEL_PROVIDER_TYPE=anthropic")
	require.Contains(t, env, "KIMI_CODE_HOME="+filepath.Join(stateDir, "kimi-code"))
	require.Empty(t, prepared.ConfigPath)
	require.NoFileExists(t, filepath.Join(stateDir, "kimi-code", "config.toml"))
	prepared.Cleanup()
}

func TestKimiACPPrepareLaunchResponsesUsesPrivateStableConfig(t *testing.T) {
	t.Setenv("KIMI_MODEL_API_KEY", "ambient-process-secret")
	t.Setenv("KIMI_MODEL_NAME", "ambient-process-model")
	stateDir := t.TempDir()
	prepared, err := prepareKimiACPLaunch(drivercontract.LaunchConfig{
		WorkDir:         t.TempDir(),
		RuntimeStateDir: stateDir,
		PermissionMode:  types.AgentPermissionModeFullAccess,
		Model:           "agentteam/gpt-5",
		Provider: &drivercontract.ProviderConfig{
			Name: "AgentTeam", BaseURL: "https://llm.example.com/v1", APIKey: "sk-secret",
			Headers: map[string]string{"X-AgentTeam": "1"}, WireAPIs: []string{provider.WireAPIResponses},
		},
	})
	require.NoError(t, err)
	require.Equal(t, filepath.Join(stateDir, "kimi-code", "config.toml"), prepared.ConfigPath)
	data, err := os.ReadFile(prepared.ConfigPath)
	require.NoError(t, err)
	var config map[string]any
	require.NoError(t, toml.Unmarshal(data, &config))
	require.Contains(t, string(data), "openai_responses")
	require.Contains(t, string(data), "sk-secret")
	require.Contains(t, string(data), "X-AgentTeam")
	require.NotContains(t, strings.Join(prepared.Env, "\n"), "KIMI_MODEL_API_KEY=")
	require.NotContains(t, strings.Join(prepared.Env, "\n"), "KIMI_MODEL_NAME=")
	require.NotContains(t, strings.Join(prepared.Env, "\n"), "ambient-process")
	sessionMarker := filepath.Join(stateDir, "kimi-code", "sessions", "keep")
	require.NoError(t, os.MkdirAll(sessionMarker, 0o700))
	prepared.Cleanup()
	require.NoFileExists(t, prepared.ConfigPath)
	require.DirExists(t, sessionMarker)
}

func TestKimiACPPrepareLaunchFailsClosedWhenNativeToolsDisabled(t *testing.T) {
	_, err := prepareKimiACPLaunch(drivercontract.LaunchConfig{
		WorkDir:            t.TempDir(),
		PermissionMode:     types.AgentPermissionModeFullAccess,
		DisableNativeTools: true,
	})
	require.ErrorContains(t, err, "cannot enforce DisableNativeTools")
}

func TestKimiACPPrepareLaunchUsesNativeConfigWithoutProvider(t *testing.T) {
	t.Setenv("KIMI_CODE_HOME", "native-kimi-home")
	t.Setenv("KIMI_MODEL_NAME", "native-kimi-model")
	stateDir := t.TempDir()
	prepared, err := prepareKimiACPLaunch(drivercontract.LaunchConfig{
		WorkDir:         t.TempDir(),
		RuntimeStateDir: stateDir,
		PermissionMode:  types.AgentPermissionModeFullAccess,
	})
	require.NoError(t, err)
	require.Equal(t, []string{"acp"}, prepared.Args)
	require.Empty(t, prepared.ConfigPath)
	require.Contains(t, strings.Join(prepared.Env, "\n"), "KIMI_CODE_HOME=native-kimi-home")
	require.Contains(t, strings.Join(prepared.Env, "\n"), "KIMI_MODEL_NAME=native-kimi-model")
	require.NoDirExists(t, filepath.Join(stateDir, "kimi-code"))
}

func TestKimiACPPrepareLaunchReplacesAmbientProviderEnvCaseInsensitively(t *testing.T) {
	t.Setenv("KIMI_MODEL_API_KEY", "ambient-process-secret")
	prepared, err := prepareKimiACPLaunch(drivercontract.LaunchConfig{
		WorkDir:         t.TempDir(),
		RuntimeStateDir: t.TempDir(),
		PermissionMode:  types.AgentPermissionModeFullAccess,
		Model:           "agentteam-model",
		Env: map[string]string{
			"kimi_model_api_key": "ambient-secret",
			"Kimi_Code_Home":     "ambient-home",
		},
		Provider: &drivercontract.ProviderConfig{
			Name: "AgentTeam", BaseURL: "https://llm.example.com/v1", APIKey: "agentteam-secret",
			WireAPIs: []string{provider.WireAIOpenAI},
		},
	})
	require.NoError(t, err)
	env := strings.Join(prepared.Env, "\n")
	require.NotContains(t, env, "ambient-secret")
	require.NotContains(t, env, "ambient-process-secret")
	require.NotContains(t, env, "ambient-home")
	require.Contains(t, env, "KIMI_MODEL_API_KEY=agentteam-secret")
	require.Equal(t, 1, strings.Count(strings.ToUpper(env), "KIMI_MODEL_API_KEY="))
}

func TestKimiACPProviderConfigCleanupDoesNotDeleteNewLaunchConfig(t *testing.T) {
	stateDir := t.TempDir()
	config := drivercontract.LaunchConfig{
		WorkDir:         t.TempDir(),
		RuntimeStateDir: stateDir,
		PermissionMode:  types.AgentPermissionModeFullAccess,
		Model:           "agentteam/gpt-5",
		Provider: &drivercontract.ProviderConfig{
			Name: "AgentTeam", BaseURL: "https://llm.example.com/v1", APIKey: "first-secret",
			Headers: map[string]string{"X-AgentTeam": "1"}, WireAPIs: []string{provider.WireAPIResponses},
		},
	}
	first, err := prepareKimiACPLaunch(config)
	require.NoError(t, err)
	config.Provider.APIKey = "second-secret"
	second, err := prepareKimiACPLaunch(config)
	require.NoError(t, err)

	first.Cleanup()
	require.FileExists(t, second.ConfigPath)
	data, err := os.ReadFile(second.ConfigPath)
	require.NoError(t, err)
	require.Contains(t, string(data), "second-secret")
	second.Cleanup()
	require.NoFileExists(t, second.ConfigPath)
}

func TestKimiACPMCPServersUseACPHeaderAndEnvArrays(t *testing.T) {
	servers, err := kimiACPMCPServers([]types.NamedMCPConfig{
		{Name: "agentteam", URL: "https://agentteam.example/mcp", Headers: map[string]string{"Z": "2", "Authorization": "Bearer token"}},
		{Name: "local", Transport: types.MCPTransportStdio, Command: "agentteam-mcp", Args: []string{"serve"}, Env: map[string]string{"B": "2", "A": "1"}},
	})
	require.NoError(t, err)
	data, err := json.Marshal(servers)
	require.NoError(t, err)
	require.JSONEq(t, `[
		{"type":"http","name":"agentteam","url":"https://agentteam.example/mcp","headers":[{"name":"Authorization","value":"Bearer token"},{"name":"Z","value":"2"}]},
		{"name":"local","command":"agentteam-mcp","args":["serve"],"env":[{"name":"A","value":"1"},{"name":"B","value":"2"}]}
	]`, string(data))
}

func TestKimiACPPromptReturnsAfterCompleteFrameAndFinishesAsynchronously(t *testing.T) {
	var stdin bytes.Buffer
	driver := newKimiACPTestDriver(&stdin)

	err := driver.DeliverRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "delivery-1"}, "hello"), "session-1")
	require.NoError(t, err)

	var request kimiACPEnvelope
	require.NoError(t, json.Unmarshal(bytes.TrimSpace(stdin.Bytes()), &request))
	require.Equal(t, "session/prompt", request.Method)
	start := parseKimiACPTestEvent(t, <-driver.output.lines)
	require.Equal(t, consts.ParsedEventKindThinking, start.Kind)
	require.True(t, parsedBoolMetadata(start.Metadata, codexAuthoritativeTurnStartMetadataKey))
	select {
	case line := <-driver.output.lines:
		t.Fatalf("terminal event emitted before prompt response: %s", line)
	default:
	}

	driver.routeEnvelope(kimiACPEnvelope{
		ID: request.ID, Result: json.RawMessage(`{"stopReason":"end_turn"}`),
	})
	terminal := waitKimiACPTestEvent(t, driver)
	require.Equal(t, consts.ParsedEventKindTurnEnd, terminal.Kind)
	require.Equal(t, consts.ParsedEventStopReasonEndTurn, terminal.StopReason)
	driver.closeTransport()
	select {
	case line := <-driver.output.lines:
		t.Fatalf("transport close duplicated a completed prompt terminal: %s", line)
	default:
	}
}

func TestKimiACPPromptPartialWriteIsAmbiguousAndClosesTransport(t *testing.T) {
	driver := newKimiACPTestDriver(io.Discard)
	driver.stdin = &shortWriteCloser{limit: 8}
	err := driver.DeliverRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "hello"), "session-1")
	require.ErrorIs(t, err, errKimiACPTransportAmbiguous)
	driver.mu.Lock()
	require.Nil(t, driver.stdin)
	require.False(t, driver.promptInFlight)
	driver.mu.Unlock()
}

func TestKimiACPAcceptedPromptGetsTerminalEventsWhenTransportCloses(t *testing.T) {
	var stdin bytes.Buffer
	driver := newKimiACPTestDriver(&stdin)
	require.NoError(t, driver.DeliverRuntimeInput(
		context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "hello"), "session-1",
	))

	start := waitKimiACPTestEvent(t, driver)
	require.Equal(t, consts.ParsedEventKindThinking, start.Kind)
	driver.closeTransport()

	failure := waitKimiACPTestEvent(t, driver)
	require.Equal(t, consts.ParsedEventKindError, failure.Kind)
	require.Contains(t, failure.Message, "transport closed")
	terminal := waitKimiACPTestEvent(t, driver)
	require.Equal(t, consts.ParsedEventKindTurnEnd, terminal.Kind)
	driver.mu.Lock()
	require.False(t, driver.promptInFlight)
	require.False(t, driver.promptAccepted)
	driver.mu.Unlock()
}

func TestKimiACPPromptPreflightFailureRemainsRetryable(t *testing.T) {
	driver := newKimiACPTestDriver(io.Discard)
	driver.stdin = nil
	err := driver.DeliverRuntimeInput(context.Background(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "hello"), "session-1")
	require.ErrorIs(t, err, errRuntimeInputMutationNotAttempted)
	require.False(t, isAmbiguousRuntimeDeliveryError(err))
}

func TestKimiACPOpenSessionInvalidResumeFallsBackOnce(t *testing.T) {
	driver := &KimiACPDriver{resumeCapable: true}
	var methods []string
	driver.requestFuncForTest = func(_ context.Context, method string, _ any) (json.RawMessage, error) {
		methods = append(methods, method)
		if method == "session/resume" {
			return nil, &kimiACPResponseError{
				Code:    -32000,
				Message: "session resume rejected",
				Data:    json.RawMessage(`{"code":"session_not_found"}`),
			}
		}
		return json.RawMessage(`{"sessionId":"fresh-session"}`), nil
	}
	sessionID, resumed, err := driver.openSession(context.Background(), drivercontract.LaunchConfig{SessionID: "old-session"}, t.TempDir())
	require.NoError(t, err)
	require.Equal(t, "fresh-session", sessionID)
	require.False(t, resumed)
	require.Equal(t, []string{"session/resume", "session/new"}, methods)
	require.Equal(t, &drivercontract.SessionResumeFallback{
		Reason:    "invalid_or_not_found_session",
		Operation: "session/resume",
	}, driver.sessionResumeFallback)
}

func TestKimiACPOpenSessionTransportFailureDoesNotFreshFallback(t *testing.T) {
	driver := &KimiACPDriver{resumeCapable: true}
	var methods []string
	driver.requestFuncForTest = func(_ context.Context, method string, _ any) (json.RawMessage, error) {
		methods = append(methods, method)
		if method == "session/resume" {
			return nil, context.DeadlineExceeded
		}
		return json.RawMessage(`{"sessionId":"unexpected-fresh"}`), nil
	}
	_, _, err := driver.openSession(context.Background(), drivercontract.LaunchConfig{SessionID: "old-session"}, t.TempDir())
	require.ErrorIs(t, err, context.DeadlineExceeded)
	require.Equal(t, []string{"session/resume"}, methods)
	require.Nil(t, driver.sessionResumeFallback)
}

func TestKimiACPOpenSessionOrdinaryProviderRejectionDoesNotFreshFallback(t *testing.T) {
	driver := &KimiACPDriver{resumeCapable: true}
	var methods []string
	driver.requestFuncForTest = func(_ context.Context, method string, _ any) (json.RawMessage, error) {
		methods = append(methods, method)
		if method == "session/resume" {
			return nil, &kimiACPResponseError{Code: -32000, Message: "model not found"}
		}
		return json.RawMessage(`{"sessionId":"unexpected-fresh"}`), nil
	}
	_, _, err := driver.openSession(context.Background(), drivercontract.LaunchConfig{SessionID: "old-session"}, t.TempDir())
	require.ErrorContains(t, err, "model not found")
	require.Equal(t, []string{"session/resume"}, methods)
	require.Nil(t, driver.sessionResumeFallback)
}

func TestKimiSessionInvalidAuthorityRequiresTypedSessionCodeAndOuterError(t *testing.T) {
	for _, data := range []string{
		`{"reason":"session_not_found"}`,
		`{"status":"session_not_found"}`,
		`{"error":"session_not_found"}`,
		`{"message":"session_not_found"}`,
		`{"code":"not_found"}`,
		`{"code":"invalid"}`,
	} {
		err := &kimiACPResponseError{
			Code: kimiSessionInvalidRPCCode,
			Data: json.RawMessage(data),
		}
		require.False(t, isAuthoritativeKimiACPSessionInvalidError("session/resume", err), "generic data must not authorize fallback: %s", data)
	}

	require.True(t, isAuthoritativeKimiACPSessionInvalidError("session/resume", &kimiACPResponseError{
		Code: kimiSessionInvalidRPCCode,
		Data: json.RawMessage(`{"type":"invalid_session_id"}`),
	}))
	require.False(t, isAuthoritativeKimiACPSessionInvalidError("session/resume", &kimiACPResponseError{
		Code: -32602,
		Data: json.RawMessage(`{"code":"session_not_found"}`),
	}), "a valid token under the wrong outer RPC code is not authority")

	driver := &KimiACPDriver{resumeCapable: true}
	var methods []string
	driver.requestFuncForTest = func(_ context.Context, method string, _ any) (json.RawMessage, error) {
		methods = append(methods, method)
		return nil, &kimiACPResponseError{
			Code: -32602,
			Data: json.RawMessage(`{"code":"session_not_found"}`),
		}
	}
	_, _, err := driver.openSession(context.Background(), drivercontract.LaunchConfig{SessionID: "old-session"}, t.TempDir())
	require.Error(t, err)
	require.Equal(t, []string{"session/resume"}, methods)
	require.Nil(t, driver.sessionResumeFallback)
}

func TestKimiACPFullAccessSelectsYoloMode(t *testing.T) {
	driver := &KimiACPDriver{}
	var method string
	var params map[string]any
	driver.requestFuncForTest = func(_ context.Context, gotMethod string, gotParams any) (json.RawMessage, error) {
		method = gotMethod
		params = gotParams.(map[string]any)
		return json.RawMessage(`{}`), nil
	}
	err := driver.setFullAccessMode(context.Background(), "session-1", drivercontract.LaunchConfig{
		PermissionMode: types.AgentPermissionModeFullAccess,
	})
	require.NoError(t, err)
	require.Equal(t, "session/set_mode", method)
	require.Equal(t, "yolo", params["modeId"])
}

func TestKimiACPUnexpectedPermissionRequestIsCancelled(t *testing.T) {
	var stdin bytes.Buffer
	driver := newKimiACPTestDriver(&stdin)
	driver.handleServerRequest(kimiACPEnvelope{
		ID: json.RawMessage(`"permission-1"`), Method: "session/request_permission",
		Params: json.RawMessage(`{"options":[{"optionId":"allow","kind":"allow_once"}]}`),
	})
	require.JSONEq(t, `{
		"jsonrpc":"2.0","id":"permission-1",
		"result":{"outcome":{"outcome":"cancelled"}}
	}`, strings.TrimSpace(stdin.String()))
}

func TestKimiACPUsageOnlyMapsContextOccupancy(t *testing.T) {
	driver := newKimiACPTestDriver(io.Discard)
	driver.handleACPNotification("session/update", json.RawMessage(`{
		"sessionId":"session-1",
		"update":{"sessionUpdate":"usage_update","used":123,"size":4096,"cost":{"amount":1}}
	}`))
	event := parseKimiACPTestEvent(t, <-driver.output.lines)
	require.Equal(t, consts.ParsedEventKindUsage, event.Kind)
	require.Equal(t, uint64(123), event.ActiveContextTokens)
	require.Equal(t, uint64(4096), event.ContextWindowTokens)
	require.Zero(t, event.InputTokens)
	require.Zero(t, event.OutputTokens)
}

func TestValidateKimiACPInitializeResultRequiresProtocolAndDetectsResume(t *testing.T) {
	resume, err := validateKimiACPInitializeResult(json.RawMessage(`{
		"protocolVersion":1,
		"agentCapabilities":{"sessionCapabilities":{"resume":{}}}
	}`))
	require.NoError(t, err)
	require.True(t, resume)
	_, err = validateKimiACPInitializeResult(json.RawMessage(`{"protocolVersion":2}`))
	require.ErrorContains(t, err, "unsupported ACP protocol")
}

func newKimiACPTestDriver(writer io.Writer) *KimiACPDriver {
	return &KimiACPDriver{
		stdin:         kimiACPWriteCloser{Writer: writer},
		output:        newKimiTranslatedOutput(32),
		pending:       make(map[int]chan kimiACPEnvelope),
		sessionID:     "session-1",
		toolTitles:    make(map[string]string),
		resumeCapable: true,
	}
}

func parseKimiACPTestEvent(t *testing.T, line string) translatedEvent {
	t.Helper()
	var event translatedEvent
	require.NoError(t, json.Unmarshal([]byte(line), &event))
	return event
}

func waitKimiACPTestEvent(t *testing.T, driver *KimiACPDriver) translatedEvent {
	t.Helper()
	select {
	case line := <-driver.output.lines:
		return parseKimiACPTestEvent(t, line)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for Kimi ACP event")
		return translatedEvent{}
	}
}

type shortWriteCloser struct {
	limit int
}

func (w *shortWriteCloser) Write(data []byte) (int, error) {
	if w.limit > len(data) {
		return len(data), nil
	}
	return w.limit, nil
}

func (*shortWriteCloser) Close() error { return nil }

type kimiACPWriteCloser struct {
	io.Writer
}

func (kimiACPWriteCloser) Close() error { return nil }