package wire

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

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

	var received struct {
		method        string
		path          string
		authorization string
		body          map[string]any
	}
	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		defer request.Body.Close()
		received.method = request.Method
		received.path = request.URL.Path
		received.authorization = request.Header.Get("Authorization")
		body, err := io.ReadAll(request.Body)
		require.NoError(t, err)
		require.NoError(t, json.Unmarshal(body, &received.body))
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"proxyToken":"proxy-token","expiresAt":"2026-06-01T01:00:00Z","sessionId":"lps_1","refreshAt":"2026-06-01T00:50:00Z"}`))
	}))
	defer server.Close()

	channel := &Protocol{
		serverURL:         server.URL,
		workspaceID:       "wsp_1",
		httpClient:        server.Client(),
		authToken:         "machine-token",
		agents:            map[string]*agentState{"agt_1": {Token: "agent-token"}},
		agentTokenRefresh: map[string]*agentTokenRefreshState{},
	}

	issuer := channel.LLMProxySessionIssuer("agt_1")
	session, err := issuer.IssueLLMProxySession(context.Background(), "lps_prev")
	require.NoError(t, err)

	require.Equal(t, http.MethodPost, received.method)
	require.Equal(t, "/api/v1/workspaces/wsp_1/agents/agt_1/runtime/llm-proxy-session", received.path)
	require.Equal(t, "Bearer agent-token", received.authorization)
	require.Equal(t, "lps_prev", received.body["sessionId"])
	require.Equal(t, "proxy-token", session.ProxyToken)
	require.Equal(t, "lps_1", session.SessionID)
	require.Equal(t, time.Date(2026, 6, 1, 1, 0, 0, 0, time.UTC), session.ExpiresAt)
	require.NotNil(t, session.RefreshAt)
}

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

	var received struct {
		body map[string]any
	}
	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		defer request.Body.Close()
		body, err := io.ReadAll(request.Body)
		require.NoError(t, err)
		require.NoError(t, json.Unmarshal(body, &received.body))
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"proxyToken":"maas-token","expiresAt":"2026-06-01T01:00:00Z","sessionId":"lps_1","maasBaseURL":"http://127.0.0.1:8080/maas/v1","credentialKind":"maas_capability"}`))
	}))
	defer server.Close()

	channel := &Protocol{
		serverURL:         server.URL,
		workspaceID:       "wsp_1",
		httpClient:        server.Client(),
		authToken:         "machine-token",
		agents:            map[string]*agentState{"agt_1": {Token: "agent-token"}},
		agentTokenRefresh: map[string]*agentTokenRefreshState{},
	}

	issuer := channel.LLMProxySessionIssuer("agt_1", agent.LLMCredentialKindMaaSCapability)
	session, err := issuer.IssueLLMProxySession(context.Background(), "")
	require.NoError(t, err)

	require.Empty(t, received.body["sessionId"])
	require.Equal(t, agent.LLMCredentialKindMaaSCapability, received.body["credentialKind"])
	require.Equal(t, "maas-token", session.ProxyToken)
	require.Equal(t, "http://127.0.0.1:8080/maas/v1", session.MaasBaseURL)
	require.Equal(t, agent.LLMCredentialKindMaaSCapability, session.CredentialKind)
}