package agent

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

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

// lockedStringBuilder is a goroutine-safe strings.Builder for slog sinks
// shared by concurrent bridge goroutines.
type lockedStringBuilder struct {
	mu sync.Mutex
	b  strings.Builder
}

func (s *lockedStringBuilder) Write(p []byte) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.b.Write(p)
}

func (s *lockedStringBuilder) String() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.b.String()
}

type recordingLLMProxySessionIssuer struct {
	mu             sync.Mutex
	issueSession   AgentTeamLLMProxySession
	refreshSession AgentTeamLLMProxySession
	issueCalls     []string
	refreshCalls   []string
	issueErr       error
	refreshErr     error
	issueErrs      []error
	refreshErrs    []error
}

func (issuer *recordingLLMProxySessionIssuer) IssueLLMProxySession(_ context.Context, selectorSessionID string) (AgentTeamLLMProxySession, error) {
	issuer.mu.Lock()
	defer issuer.mu.Unlock()
	issuer.issueCalls = append(issuer.issueCalls, strings.TrimSpace(selectorSessionID))
	return issuer.issueSession, issuer.nextIssueErr()
}

func (issuer *recordingLLMProxySessionIssuer) RefreshLLMProxySession(_ context.Context, selectorSessionID string) (AgentTeamLLMProxySession, error) {
	issuer.mu.Lock()
	defer issuer.mu.Unlock()
	issuer.refreshCalls = append(issuer.refreshCalls, strings.TrimSpace(selectorSessionID))
	return issuer.refreshSession, issuer.nextRefreshErr()
}

func (issuer *recordingLLMProxySessionIssuer) nextIssueErr() error {
	if len(issuer.issueErrs) == 0 {
		return issuer.issueErr
	}
	err := issuer.issueErrs[0]
	issuer.issueErrs = issuer.issueErrs[1:]
	return err
}

func (issuer *recordingLLMProxySessionIssuer) nextRefreshErr() error {
	if len(issuer.refreshErrs) == 0 {
		return issuer.refreshErr
	}
	err := issuer.refreshErrs[0]
	issuer.refreshErrs = issuer.refreshErrs[1:]
	return err
}

func TestAgentTeamMaaSLLMBridgeInjectsMaaSCapabilityAndIdempotencyKey(t *testing.T) {
	t.Parallel()

	var received struct {
		authorization string
		idempotency   string
		path          string
	}
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		received.authorization = request.Header.Get("Authorization")
		received.idempotency = request.Header.Get("Idempotency-Key")
		received.path = request.URL.Path
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	issuer := &recordingLLMProxySessionIssuer{
		issueSession: AgentTeamLLMProxySession{
			SessionID:  "lps_1",
			ProxyToken: "maas-capability-1",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bridge, err := StartAgentTeamMaaSLLMBridge(ctx, AgentTeamMaaSLLMBridgeConfig{
		UpstreamBaseURL: upstream.URL + "/maas/v1",
		Issuer:          issuer,
	})
	require.NoError(t, err)
	defer func() { require.NoError(t, bridge.Close()) }()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, bridge.BaseURL()+"/messages", strings.NewReader(`{"model":"hosted/deepseek_chat"}`))
	require.NoError(t, err)
	req.Header.Set("X-Api-Key", bridge.APIKey())
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	_, err = io.ReadAll(resp.Body)
	require.NoError(t, err)

	require.Equal(t, "Bearer maas-capability-1", received.authorization)
	require.NotEmpty(t, received.idempotency)
	require.Equal(t, "/maas/v1/messages", received.path)
	require.Equal(t, []string{""}, issuer.issueCalls)
	require.Empty(t, issuer.refreshCalls)
}

func TestAgentTeamMaaSLLMBridgeForwardsDirectMaaSPath(t *testing.T) {
	t.Parallel()

	var received struct {
		authorization     string
		idempotency       string
		path              string
		body              string
		agentAuthToken    string
		xAgentAuthToken   string
		xAPIKey           string
		maasRequestID     string
		maasUsageEventID  string
		maasReservationID string
	}
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		body, err := io.ReadAll(request.Body)
		if err != nil {
			writer.WriteHeader(http.StatusInternalServerError)
			return
		}
		received.authorization = request.Header.Get("Authorization")
		received.idempotency = request.Header.Get("Idempotency-Key")
		received.path = request.URL.Path
		received.body = string(body)
		received.agentAuthToken = request.Header.Get("AgentAuthToken")
		received.xAgentAuthToken = request.Header.Get("X-Agent-Auth-Token")
		received.xAPIKey = request.Header.Get("X-Api-Key")
		received.maasRequestID = request.Header.Get("X-AgentTeam-Maas-Request-Id")
		received.maasUsageEventID = request.Header.Get("X-AgentTeam-Maas-Usage-Event-Id")
		received.maasReservationID = request.Header.Get("X-AgentTeam-Maas-Reservation-Id")
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	issuer := &recordingLLMProxySessionIssuer{
		issueSession: AgentTeamLLMProxySession{
			SessionID:      "lps_1",
			ProxyToken:     "maas-capability-1",
			ExpiresAt:      time.Now().UTC().Add(time.Minute),
			MaasBaseURL:    upstream.URL + "/maas/v1",
			CredentialKind: LLMCredentialKindMaaSCapability,
		},
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bridge, err := StartAgentTeamMaaSLLMBridge(ctx, AgentTeamMaaSLLMBridgeConfig{
		UpstreamBaseURL: upstream.URL + "/maas/v1",
		Issuer:          issuer,
	})
	require.NoError(t, err)
	defer func() { require.NoError(t, bridge.Close()) }()

	const requestBody = `{"model":"hosted/deepseek_chat","output_config":{"effort":"high"},"context_management":{"auto_clear":true}}`
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, bridge.BaseURL()+"/messages", strings.NewReader(requestBody))
	require.NoError(t, err)
	req.Header.Set("X-Api-Key", bridge.APIKey())
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("AgentAuthToken", "legacy-agent-token")
	req.Header.Set("X-Agent-Auth-Token", "legacy-agent-token")
	req.Header.Set("X-AgentTeam-Maas-Request-Id", "spoofed-request")
	req.Header.Set("X-AgentTeam-Maas-Usage-Event-Id", "spoofed-usage")
	req.Header.Set("X-AgentTeam-Maas-Reservation-Id", "spoofed-reservation")

	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	_, err = io.ReadAll(resp.Body)
	require.NoError(t, err)

	require.True(t, strings.HasSuffix(bridge.BaseURL(), "/maas/v1"))
	require.Equal(t, "Bearer maas-capability-1", received.authorization)
	require.NotEmpty(t, received.idempotency)
	require.Equal(t, "/maas/v1/messages", received.path)
	require.Equal(t, requestBody, received.body)
	require.Empty(t, received.agentAuthToken)
	require.Empty(t, received.xAgentAuthToken)
	require.Empty(t, received.xAPIKey)
	require.Empty(t, received.maasRequestID)
	require.Empty(t, received.maasUsageEventID)
	require.Empty(t, received.maasReservationID)
	require.Equal(t, []string{""}, issuer.issueCalls)
	require.Empty(t, issuer.refreshCalls)
}

func TestAgentTeamMaaSLLMBridgeSkipsIdempotencyKeyForCountTokens(t *testing.T) {
	t.Parallel()

	var received struct {
		authorization string
		idempotency   string
		path          string
	}
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		received.authorization = request.Header.Get("Authorization")
		received.idempotency = request.Header.Get("Idempotency-Key")
		received.path = request.URL.Path
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	issuer := &recordingLLMProxySessionIssuer{
		issueSession: AgentTeamLLMProxySession{
			SessionID:  "lps_1",
			ProxyToken: "maas-capability-1",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bridge, err := StartAgentTeamMaaSLLMBridge(ctx, AgentTeamMaaSLLMBridgeConfig{
		UpstreamBaseURL: upstream.URL + "/maas/v1",
		Issuer:          issuer,
	})
	require.NoError(t, err)
	defer func() { require.NoError(t, bridge.Close()) }()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, bridge.BaseURL()+"/messages/count_tokens", strings.NewReader(`{"model":"hosted/deepseek_chat"}`))
	require.NoError(t, err)
	req.Header.Set("X-Api-Key", bridge.APIKey())
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	_, err = io.ReadAll(resp.Body)
	require.NoError(t, err)

	require.Equal(t, "Bearer maas-capability-1", received.authorization)
	require.Empty(t, received.idempotency)
	require.Equal(t, "/maas/v1/messages/count_tokens", received.path)
	require.Equal(t, []string{""}, issuer.issueCalls)
	require.Empty(t, issuer.refreshCalls)
}

func TestAgentTeamMaaSLLMBridgeRefreshesMaaSCapabilityAfterUnauthorized(t *testing.T) {
	t.Parallel()

	var mu sync.Mutex
	var authHeaders []string
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		mu.Lock()
		authHeaders = append(authHeaders, request.Header.Get("Authorization"))
		call := len(authHeaders)
		mu.Unlock()

		if call == 1 {
			writer.WriteHeader(http.StatusUnauthorized)
			_, _ = writer.Write([]byte(`{"error":"expired"}`))
			return
		}
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	issuer := &recordingLLMProxySessionIssuer{
		issueSession: AgentTeamLLMProxySession{
			SessionID:  "lps_1",
			ProxyToken: "maas-capability-1",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
		refreshSession: AgentTeamLLMProxySession{
			SessionID:  "lps_2",
			ProxyToken: "maas-capability-2",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bridge, err := StartAgentTeamMaaSLLMBridge(ctx, AgentTeamMaaSLLMBridgeConfig{
		UpstreamBaseURL: upstream.URL + "/maas/v1",
		Issuer:          issuer,
	})
	require.NoError(t, err)
	defer func() { require.NoError(t, bridge.Close()) }()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, bridge.BaseURL()+"/messages", strings.NewReader(`{"model":"hosted/deepseek_chat"}`))
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+bridge.APIKey())
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	require.NoError(t, err)

	require.Equal(t, http.StatusOK, resp.StatusCode)
	require.JSONEq(t, `{"ok":true}`, string(body))
	require.Equal(t, []string{""}, issuer.issueCalls)
	require.Equal(t, []string{"lps_1"}, issuer.refreshCalls)
	require.Equal(t, []string{"Bearer maas-capability-1", "Bearer maas-capability-2"}, authHeaders)
}

func TestAgentTeamMaaSLLMBridgePreservesUpstreamStreamAbort(t *testing.T) {
	t.Parallel()

	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		writer.Header().Set("Content-Type", "text/event-stream")
		_, _ = writer.Write([]byte("data: partial\n\n"))
		writer.(http.Flusher).Flush()
		panic(http.ErrAbortHandler)
	}))
	defer upstream.Close()

	issuer := &recordingLLMProxySessionIssuer{
		issueSession: AgentTeamLLMProxySession{
			SessionID:  "lps_1",
			ProxyToken: "maas-capability-1",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bridge, err := StartAgentTeamMaaSLLMBridge(ctx, AgentTeamMaaSLLMBridgeConfig{
		UpstreamBaseURL: upstream.URL + "/maas/v1",
		Issuer:          issuer,
	})
	require.NoError(t, err)
	defer func() { require.NoError(t, bridge.Close()) }()

	request, err := http.NewRequestWithContext(ctx, http.MethodPost, bridge.BaseURL()+"/messages", strings.NewReader(`{"model":"hosted/deepseek_chat"}`))
	require.NoError(t, err)
	request.Header.Set("X-Api-Key", bridge.APIKey())

	response, err := http.DefaultClient.Do(request)
	require.NoError(t, err)
	defer response.Body.Close()
	body, readErr := io.ReadAll(response.Body)

	require.Equal(t, http.StatusOK, response.StatusCode)
	require.Equal(t, "data: partial\n\n", string(body))
	require.Error(t, readErr)
}

func TestAgentTeamMaaSLLMBridgeFlushesStreamChunks(t *testing.T) {
	t.Parallel()

	releaseUpstream := make(chan struct{})
	var releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(releaseUpstream) }) }
	defer release()
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "text/event-stream")
		_, _ = writer.Write([]byte("data: first\n\n"))
		writer.(http.Flusher).Flush()
		select {
		case <-releaseUpstream:
			_, _ = writer.Write([]byte("data: [DONE]\n\n"))
		case <-request.Context().Done():
		}
	}))
	defer upstream.Close()

	issuer := &recordingLLMProxySessionIssuer{
		issueSession: AgentTeamLLMProxySession{
			SessionID:  "lps_1",
			ProxyToken: "maas-capability-1",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bridge, err := StartAgentTeamMaaSLLMBridge(ctx, AgentTeamMaaSLLMBridgeConfig{
		UpstreamBaseURL: upstream.URL + "/maas/v1",
		Issuer:          issuer,
	})
	require.NoError(t, err)
	defer func() { require.NoError(t, bridge.Close()) }()

	request, err := http.NewRequestWithContext(ctx, http.MethodPost, bridge.BaseURL()+"/messages", strings.NewReader(`{"model":"hosted/deepseek_chat"}`))
	require.NoError(t, err)
	request.Header.Set("X-Api-Key", bridge.APIKey())
	client := &http.Client{Timeout: 5 * time.Second}

	response, err := client.Do(request)
	require.NoError(t, err)
	defer response.Body.Close()
	release()
	body, err := io.ReadAll(response.Body)
	require.NoError(t, err)

	require.Equal(t, "data: first\n\ndata: [DONE]\n\n", string(body))
}

func TestAgentTeamMaaSLLMBridgeRuntimeDebugScrubsLogs(t *testing.T) {
	// The bridge logs from its proxy-handler goroutines, so the shared
	// default-handler sink must be goroutine-safe; a bare strings.Builder
	// races under -race when concurrent requests log in parallel.
	var logBuffer lockedStringBuilder
	logger := slog.New(slog.NewTextHandler(&logBuffer, &slog.HandlerOptions{Level: slog.LevelInfo}))
	previousLogger := slog.Default()
	slog.SetDefault(logger)
	t.Cleanup(func() { slog.SetDefault(previousLogger) })

	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"ok":true}`))
	}))
	defer upstream.Close()

	issuer := &recordingLLMProxySessionIssuer{
		issueSession: AgentTeamLLMProxySession{
			SessionID:  "lps_1",
			ProxyToken: "maas-capability-secret",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bridge, err := StartAgentTeamMaaSLLMBridge(ctx, AgentTeamMaaSLLMBridgeConfig{
		UpstreamBaseURL:     upstream.URL + "/maas/v1?token=secret",
		Issuer:              issuer,
		RuntimeDebugEnabled: true,
		AgentID:             "agt_debug",
		ServerAgentID:       "server_agt_debug",
	})
	require.NoError(t, err)
	defer func() { require.NoError(t, bridge.Close()) }()

	request, err := http.NewRequestWithContext(ctx, http.MethodPost, bridge.BaseURL()+"/messages?api_key=secret", strings.NewReader(`{"model":"hosted/deepseek_chat","prompt":"should not be logged"}`))
	require.NoError(t, err)
	request.Header.Set("X-Api-Key", bridge.APIKey())
	response, err := http.DefaultClient.Do(request)
	require.NoError(t, err)
	defer response.Body.Close()
	_, err = io.ReadAll(response.Body)
	require.NoError(t, err)

	logs := logBuffer.String()
	require.Contains(t, logs, "agentteam maas llm bridge upstream request sending")
	require.Contains(t, logs, "query_present=true")
	require.Contains(t, logs, "body_bytes=")
	require.NotContains(t, logs, "maas-capability-secret")
	require.NotContains(t, logs, bridge.APIKey())
	require.NotContains(t, logs, "token=secret")
	require.NotContains(t, logs, "api_key=secret")
	require.NotContains(t, logs, "should not be logged")
}

func TestAgentTeamMaaSLLMBridgeRetriesIssueAfterRuntimeSessionNotReady(t *testing.T) {
	t.Parallel()

	issuer := &recordingLLMProxySessionIssuer{
		issueSession: AgentTeamLLMProxySession{
			SessionID:  "lps_1",
			ProxyToken: "maas-capability-1",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
		issueErrs: []error{errors.New("agentteam api post: HTTP 409: agent runtime session is not ready")},
	}
	var published []string
	var waits int
	bridge := &AgentTeamMaaSLLMBridge{
		issuer:           issuer,
		runtimeSessionID: "ars_runtime",
		runtimeSessionPublisher: func(_ context.Context, runtimeSessionID string) error {
			published = append(published, strings.TrimSpace(runtimeSessionID))
			return nil
		},
		runtimeSessionNotReadyBackoff: func(context.Context) error {
			waits++
			return nil
		},
	}

	token, err := bridge.ensureMaaSProxyToken(context.Background(), false)
	require.NoError(t, err)
	require.Equal(t, "maas-capability-1", token)
	require.Equal(t, []string{"", ""}, issuer.issueCalls)
	require.Empty(t, issuer.refreshCalls)
	require.Equal(t, []string{"ars_runtime", "ars_runtime"}, published)
	require.Equal(t, 1, waits)
}

func TestAgentTeamMaaSLLMBridgeRetriesRefreshAfterRuntimeSessionNotReady(t *testing.T) {
	t.Parallel()

	issuer := &recordingLLMProxySessionIssuer{
		refreshSession: AgentTeamLLMProxySession{
			SessionID:  "lps_2",
			ProxyToken: "maas-capability-2",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
		refreshErrs: []error{errors.New("agentteam api post: HTTP 409: agent_runtime_session_not_ready")},
	}
	var published []string
	var waits int
	bridge := &AgentTeamMaaSLLMBridge{
		issuer: issuer,
		session: AgentTeamLLMProxySession{
			SessionID:  "lps_1",
			ProxyToken: "maas-capability-1",
			ExpiresAt:  time.Now().UTC().Add(time.Minute),
		},
		runtimeSessionID: "ars_runtime",
		runtimeSessionPublisher: func(_ context.Context, runtimeSessionID string) error {
			published = append(published, strings.TrimSpace(runtimeSessionID))
			return nil
		},
		runtimeSessionNotReadyBackoff: func(context.Context) error {
			waits++
			return nil
		},
	}

	token, err := bridge.ensureMaaSProxyToken(context.Background(), true)
	require.NoError(t, err)
	require.Equal(t, "maas-capability-2", token)
	require.Empty(t, issuer.issueCalls)
	require.Equal(t, []string{"lps_1", "lps_1"}, issuer.refreshCalls)
	require.Equal(t, []string{"ars_runtime", "ars_runtime"}, published)
	require.Equal(t, 1, waits)
}

func TestAgentTeamMaaSLLMBridgeUsesNoHTTPClientTimeout(t *testing.T) {
	t.Parallel()

	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.WriteHeader(http.StatusNoContent)
	}))
	defer upstream.Close()

	issuer := &recordingLLMProxySessionIssuer{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bridge, err := StartAgentTeamMaaSLLMBridge(ctx, AgentTeamMaaSLLMBridgeConfig{
		UpstreamBaseURL: upstream.URL + "/maas/v1",
		Issuer:          issuer,
	})
	require.NoError(t, err)
	defer func() { require.NoError(t, bridge.Close()) }()

	require.NotNil(t, bridge.httpClient)
	require.Zero(t, bridge.httpClient.Timeout)
}