package agent

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

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
)

type fakeLLMProxySessionStore struct {
	maasIssueSelector   string
	maasRefreshSelector string
	maasResult          agent.LLMProxySession
	maasErr             error
}

func (fake *fakeLLMProxySessionStore) IssueMaaSLLMProxySession(_ context.Context, _ agent.AuthenticatedAgent, selectorSessionID string) (agent.LLMProxySession, error) {
	fake.maasIssueSelector = selectorSessionID
	return fake.maasResult, fake.maasErr
}

func (fake *fakeLLMProxySessionStore) RefreshMaaSLLMProxySession(_ context.Context, _ agent.AuthenticatedAgent, selectorSessionID string) (agent.LLMProxySession, error) {
	fake.maasRefreshSelector = selectorSessionID
	return fake.maasResult, fake.maasErr
}

func TestLLMProxySessionHandlerRefreshMapsConflict(t *testing.T) {
	gin.SetMode(gin.TestMode)

	store := &fakeLLMProxySessionStore{
		maasErr: storeerr.Conflict("llm proxy session selector did not match current session"),
	}
	handler := &LLMProxySessionHandler{
		auth:     fakeAgentAuthenticator{agent: agent.AuthenticatedAgent{AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test"}},
		sessions: store,
	}

	router := gin.New()
	router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/llm-proxy-session/refresh", handler.Refresh)

	body, err := json.Marshal(map[string]string{"credentialKind": LLMProxyCredentialKindMaaSCapability})
	require.NoError(t, err)
	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_test/agents/agt_0000000000test/runtime/llm-proxy-session/refresh", bytes.NewReader(body))
	request.Header.Set("Authorization", "Bearer test-token")
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()

	router.ServeHTTP(response, request)

	require.Equal(t, http.StatusConflict, response.Code)
	require.Empty(t, store.maasRefreshSelector)
	require.JSONEq(t, `{"error":"llm proxy session selector did not match current session"}`, response.Body.String())
}

func TestLLMProxySessionHandlerIssueMaaSCapability(t *testing.T) {
	gin.SetMode(gin.TestMode)

	refreshAt := time.Date(2026, 6, 1, 10, 10, 0, 0, time.UTC)
	expiresAt := time.Date(2026, 6, 1, 10, 15, 0, 0, time.UTC)
	store := &fakeLLMProxySessionStore{
		maasResult: agent.LLMProxySession{
			SessionID:           "lps_maas",
			ProxyToken:          "maas.jwt.token",
			RefreshAt:           &refreshAt,
			ExpiresAt:           expiresAt,
			SupersedesSessionID: "lps_prev",
			MaasBaseURL:         "http://127.0.0.1:8080/maas/v1",
			CredentialKind:      LLMProxyCredentialKindMaaSCapability,
		},
	}
	handler := &LLMProxySessionHandler{
		auth:     fakeAgentAuthenticator{agent: agent.AuthenticatedAgent{AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test"}},
		sessions: store,
	}

	router := gin.New()
	router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/llm-proxy-session", handler.Issue)

	body, err := json.Marshal(map[string]string{
		"sessionId":      "lps_prev",
		"credentialKind": LLMProxyCredentialKindMaaSCapability,
	})
	require.NoError(t, err)
	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_test/agents/agt_0000000000test/runtime/llm-proxy-session", bytes.NewReader(body))
	request.Header.Set("Authorization", "Bearer test-token")
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()

	router.ServeHTTP(response, request)

	require.Equal(t, http.StatusOK, response.Code)
	require.Equal(t, "lps_prev", store.maasIssueSelector)
	require.JSONEq(t, `{
		"proxyToken":"maas.jwt.token",
		"expiresAt":"2026-06-01T10:15:00Z",
		"sessionId":"lps_maas",
		"refreshAt":"2026-06-01T10:10:00Z",
		"supersedesSessionId":"lps_prev",
		"maasBaseURL":"http://127.0.0.1:8080/maas/v1",
		"credentialKind":"maas_capability"
	}`, response.Body.String())
}

func TestLLMProxySessionHandlerIssueMapsRuntimeSessionNotReadyCode(t *testing.T) {
	gin.SetMode(gin.TestMode)

	store := &fakeLLMProxySessionStore{
		maasErr: maasRuntimeSessionNotReadyError{},
	}
	handler := &LLMProxySessionHandler{
		auth:     fakeAgentAuthenticator{agent: agent.AuthenticatedAgent{AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test"}},
		sessions: store,
	}

	router := gin.New()
	router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/llm-proxy-session", handler.Issue)

	body, err := json.Marshal(map[string]string{"credentialKind": LLMProxyCredentialKindMaaSCapability})
	require.NoError(t, err)
	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_test/agents/agt_0000000000test/runtime/llm-proxy-session", bytes.NewReader(body))
	request.Header.Set("Authorization", "Bearer test-token")
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()

	router.ServeHTTP(response, request)

	require.Equal(t, http.StatusConflict, response.Code)
	require.JSONEq(t, `{"error":"agent runtime session is not ready","code":"agent_runtime_session_not_ready"}`, response.Body.String())
}

func TestLLMProxySessionHandlerRejectsUnsupportedCredentialKind(t *testing.T) {
	gin.SetMode(gin.TestMode)

	handler := &LLMProxySessionHandler{
		auth:     fakeAgentAuthenticator{agent: agent.AuthenticatedAgent{AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test"}},
		sessions: &fakeLLMProxySessionStore{},
	}

	router := gin.New()
	router.POST("/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/llm-proxy-session", handler.Issue)

	body, err := json.Marshal(map[string]string{"credentialKind": "agentAuthToken"})
	require.NoError(t, err)
	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_test/agents/agt_0000000000test/runtime/llm-proxy-session", bytes.NewReader(body))
	request.Header.Set("Authorization", "Bearer test-token")
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()

	router.ServeHTTP(response, request)

	require.Equal(t, http.StatusBadRequest, response.Code)
	require.JSONEq(t, `{"error":"unsupported llm proxy credentialKind"}`, response.Body.String())
}