package llm

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

	"github.com/tidbcloud/AgentTeam/internal/apr/agentteambridge"
	"github.com/tidbcloud/AgentTeam/internal/apr/runtime"
	"github.com/tidbcloud/AgentTeam/internal/store/organization"

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

type codexRegistryTestLoader struct {
	providers []organization.OrganizationProviderView
	routes    []organization.OrganizationProviderRouteView
	groups    []organization.OrganizationProviderRouteGroupView
	tokens    map[string][]organization.OrganizationProviderTokenSecret
}

func (loader codexRegistryTestLoader) ListActiveOrganizationProvidersForLLM(context.Context) ([]organization.OrganizationProviderView, error) {
	return append([]organization.OrganizationProviderView(nil), loader.providers...), nil
}

func (loader codexRegistryTestLoader) ListActiveOrganizationProviderRoutesForLLM(context.Context) ([]organization.OrganizationProviderRouteView, error) {
	return append([]organization.OrganizationProviderRouteView(nil), loader.routes...), nil
}

func (loader codexRegistryTestLoader) ListActiveOrganizationProviderRouteGroupsForLLM(context.Context) ([]organization.OrganizationProviderRouteGroupView, error) {
	return append([]organization.OrganizationProviderRouteGroupView(nil), loader.groups...), nil
}

func (loader codexRegistryTestLoader) LoadActiveOrganizationProviderTokens(_ context.Context, organizationID string, providerID string) ([]organization.OrganizationProviderTokenSecret, error) {
	key := organizationID + "/" + providerID
	return append([]organization.OrganizationProviderTokenSecret(nil), loader.tokens[key]...), nil
}

func (codexRegistryTestLoader) UpdateOrganizationProviderTokenSecret(context.Context, string, string, string, organization.UpdateOrganizationProviderTokenSecretParams) error {
	return nil
}

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

	var receivedPath string
	var receivedAPIKey string
	var receivedBody map[string]any
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		receivedPath = request.URL.Path
		receivedAPIKey = request.Header.Get("X-Api-Key")
		body, _ := io.ReadAll(request.Body)
		_ = json.Unmarshal(body, &receivedBody)
		writer.Header().Set("Content-Type", "text/event-stream")
		_, _ = writer.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"created_at\":1735689600,\"model\":\"gpt-5.3-codex\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"hello from codex\"}]}],\"usage\":{\"input_tokens\":11,\"output_tokens\":7,\"total_tokens\":18}}}\n\n"))
		_, _ = writer.Write([]byte("data: [DONE]\n\n"))
	}))
	defer upstream.Close()

	dispatcher := runtime.NewDispatcher()
	loader := codexRegistryTestLoader{
		providers: []organization.OrganizationProviderView{{
			ID:             "opvd_codex",
			OrganizationID: "org_1",
			Name:           "Codex",
			DriverKind:     "codex",
			BaseURL:        upstream.URL,
			Capabilities: organization.OrganizationProviderCapabilities{
				OpenAICompatible:    true,
				ResponsesCompatible: true,
			},
			Status: "active",
		}},
		routes: []organization.OrganizationProviderRouteView{{
			ID:             "oprt_codex",
			OrganizationID: "org_1",
			Vendor:         "openai",
			DisplayName:    "OpenAI",
			Status:         "active",
		}},
		groups: []organization.OrganizationProviderRouteGroupView{{
			ID:             "ogrp_codex",
			OrganizationID: "org_1",
			RouteID:        "oprt_codex",
			Models:         []string{"gpt-5.3-codex"},
			ProviderRefs:   []string{"opvd_codex"},
			Status:         "active",
		}},
		tokens: map[string][]organization.OrganizationProviderTokenSecret{
			"org_1/opvd_codex": {{
				ID:    "optk_codex",
				Label: "primary",
				Config: organization.OrganizationProviderTokenConfig{
					Role: "primary",
					Mode: "api_key",
				},
				Secret: organization.OrganizationProviderTokenMaterial{
					APIKey: "sk-codex",
				},
			}},
		},
	}
	_, cleanups, _, err := agentteambridge.ApplyOrganizationProviders(context.Background(), dispatcher, loader, upstream.Client(), nil)
	require.NoError(t, err)
	defer func() {
		for i := len(cleanups) - 1; i >= 0; i-- {
			if cleanups[i] != nil {
				cleanups[i]()
			}
		}
	}()

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterRoutes(router, Dependencies{APRDispatcher: dispatcher, ProxyKeyAuthenticator: fakeProxyKeyAuthenticator{}})

	request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/llm/v1/chat/completions", bytes.NewBufferString(`{"model":"openai/gpt-5.3-codex","messages":[{"role":"user","content":"hi"}]}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Authorization", "Bearer opvd_codex")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
	require.Equal(t, "/responses", receivedPath)
	require.Equal(t, "sk-codex", receivedAPIKey)
	require.Equal(t, "gpt-5.3-codex", receivedBody["model"])

	var payload struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &payload))
	require.Len(t, payload.Choices, 1)
	require.Equal(t, "hello from codex", payload.Choices[0].Message.Content)
}