package workspace

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

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

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/api/dto"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	"github.com/tidbcloud/AgentTeam/internal/store/identity"
	"github.com/tidbcloud/AgentTeam/internal/store/organization"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)

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

	var output bytes.Buffer
	logger := slog.New(logging.WrapHandler(slog.NewJSONHandler(&output, nil))).With(slog.String("domain", "workspace"))
	ctx := logging.WithRequestID(context.Background(), "req_models")
	logHostedAvailabilityFailure(ctx, logger, "org_1", "wsp_1", errors.New("secret upstream response"))

	var entry map[string]any
	require.NoError(t, json.Unmarshal(bytes.TrimSpace(output.Bytes()), &entry))
	require.Equal(t, "workspace", entry["domain"])
	require.Equal(t, "req_models", entry["request_id"])
	require.Equal(t, "org_1", entry["organization_id"])
	require.Equal(t, "wsp_1", entry["workspace_id"])
	require.Equal(t, "hosted_availability_unavailable", entry["error_class"])
	require.NotContains(t, output.String(), "secret upstream response")
	require.NotContains(t, entry, "organizationID")
	require.NotContains(t, entry, "workspaceID")
}

type fakeHostedAvailabilityClient struct {
	organizationID string
	availability   HostedAvailability
	err            error
}

func (fake *fakeHostedAvailabilityClient) LoadHostedAvailability(_ context.Context, organizationID string) (HostedAvailability, error) {
	fake.organizationID = organizationID
	if fake.err != nil {
		return HostedAvailability{}, fake.err
	}
	return fake.availability, nil
}

type fakeHostedPricingClient struct {
	organizationID string
	pricing        HostedPricing
	err            error
}

func (fake *fakeHostedPricingClient) LoadHostedPricing(_ context.Context, organizationID string) (HostedPricing, error) {
	fake.organizationID = organizationID
	if fake.err != nil {
		return HostedPricing{}, fake.err
	}
	return fake.pricing, nil
}

func int64Ptr(value int64) *int64 {
	return &value
}

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

	hosted := &fakeHostedAvailabilityClient{err: errors.New("agentteam-maas unavailable")}
	handler := &WorkspaceHandler{
		organizations: selectableModelOrganizationProvider("org_owner", "show_disabled_when_unavailable", "all_active_hosted_inventory"),
		hostedModels:  hosted,
	}

	response, err := handler.SelectableModels(selectableModelRequest(t, ""))
	require.NoError(t, err)
	require.Equal(t, http.StatusOK, response.Status)

	body := response.Body.(dto.SelectableModelsResponse)
	require.Equal(t, hostedAvailabilityStatusUnavailable, body.HostedAvailability.Status)
	require.Equal(t, []string{hostedReasonTemporarilyUnavailable}, body.HostedAvailability.Reasons)
	require.Equal(t, "org_1", hosted.organizationID)
	require.Len(t, body.Models, 1)
	require.Equal(t, selectableSourceBYOK, body.Models[0].Source)
	require.Equal(t, "claude-sonnet-4.6", body.Models[0].ModelReference)
	require.True(t, body.Models[0].Enabled)
}

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

	pricing := &fakeHostedPricingClient{pricing: selectableHostedPricing()}
	handler := &WorkspaceHandler{
		organizations: selectableModelOrganizationProvider("org_owner", "show_disabled_when_unavailable", "workspace_not_authorized"),
		hostedModels:  &fakeHostedAvailabilityClient{availability: selectableHostedAvailability(true, nil)},
		hostedPricing: pricing,
	}

	response, err := handler.SelectableModels(selectableModelRequest(t, ""))
	require.NoError(t, err)
	body := response.Body.(dto.SelectableModelsResponse)

	require.Equal(t, hostedAvailabilityStatusAvailable, body.HostedAvailability.Status)
	require.Equal(t, int64(123_456), body.HostedAvailability.CreditAtomsPerCredit)
	require.Equal(t, "org_1", pricing.organizationID)
	require.Len(t, body.Models, 2)
	hosted := requireSelectableModel(t, body.Models, selectableSourceHosted)
	require.False(t, hosted.Enabled)
	require.Equal(t, []string{hostedReasonWorkspaceNotAuthorized}, hosted.Reasons)
	require.Equal(t, []string{selectableActionManageWorkspaceAuthorization}, hosted.Actions)
	require.Equal(t, hostedMachineCapabilityMaaSDirect, hosted.RequiresMachineCapability)
	require.Equal(t, "org_grant", hosted.Hosted.AccessSource)
	require.Equal(t, []string{"messages", "responses"}, hosted.Hosted.SupportedFamilies)
	require.Equal(t, []string{"responses"}, hosted.Hosted.ExecutableWireAPIFamilies)
	require.Equal(t, []string{"tool_use"}, hosted.Hosted.MessagesFeatures)
	require.Equal(t, []string{"json_object"}, hosted.Hosted.ChatCompletionsFeatures)
	require.Equal(t, []string{"text_format"}, hosted.Hosted.ResponsesFeatures)
	require.NotNil(t, hosted.Hosted.InputCreditAtomsPerMillionTokens)
	require.NotNil(t, hosted.Hosted.OutputCreditAtomsPerMillionTokens)
	require.Equal(t, int64(2_000_000), *hosted.Hosted.InputCreditAtomsPerMillionTokens)
	require.Equal(t, int64(8_000_000), *hosted.Hosted.OutputCreditAtomsPerMillionTokens)
	require.Equal(t, "hpv_deepseek_v4_pro_discount_90", hosted.Hosted.EffectivePricingVersionID)
	require.Equal(t, "hpv_deepseek_v4_pro_2026_06_15", hosted.Hosted.BasePricingVersionID)
	require.Equal(t, "hosted_token_context_threshold_v1", hosted.Hosted.PricingFormulaVersion)
	require.Equal(t, []string{"input_tokens", "cache_read_input_tokens", "output_tokens"}, hosted.Hosted.BillableDimensions)
	require.Len(t, hosted.Hosted.ContextTiers, 1)
	require.NotNil(t, hosted.Hosted.ContextTiers[0].MaxContextTokens)
	require.Equal(t, int64(4096), *hosted.Hosted.ContextTiers[0].MaxContextTokens)
	require.Len(t, hosted.Hosted.BaseContextTiers, 1)
	require.NotNil(t, hosted.Hosted.PricingOverride)
	require.Equal(t, "90% campaign", hosted.Hosted.PricingOverride.DisplayLabel)
}

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

	handler := &WorkspaceHandler{
		organizations: selectableModelOrganizationProvider("org_owner", "hide_disabled_when_unavailable", "workspace_not_authorized"),
		hostedModels:  &fakeHostedAvailabilityClient{availability: selectableHostedAvailability(true, nil)},
	}

	response, err := handler.SelectableModels(selectableModelRequest(t, ""))
	require.NoError(t, err)
	body := response.Body.(dto.SelectableModelsResponse)

	require.Len(t, body.Models, 1)
	require.Equal(t, selectableSourceBYOK, body.Models[0].Source)
}

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

	handler := &WorkspaceHandler{
		organizations: selectableModelOrganizationProvider("org_member", "show_disabled_when_unavailable", "all_active_hosted_inventory"),
		hostedModels:  &fakeHostedAvailabilityClient{availability: selectableHostedAvailability(true, []string{"tool_use"})},
	}

	response, err := handler.SelectableModels(selectableModelRequest(t, "?features=vision,tool_use"))
	require.NoError(t, err)
	body := response.Body.(dto.SelectableModelsResponse)

	hosted := requireSelectableModel(t, body.Models, selectableSourceHosted)
	require.False(t, hosted.Enabled)
	require.Equal(t, []string{hostedReasonFeatureUnsupported}, hosted.Reasons)
	require.Equal(t, []string{selectableActionContactAdmin}, hosted.Actions)
}

func selectableModelOrganizationProvider(role string, pickerPolicy string, authorizationMode string) *fakeOrganizationProvider {
	return &fakeOrganizationProvider{
		providers: []organization.OrganizationProviderView{{
			ID:         "op_1",
			Name:       "Anthropic",
			DriverKind: "openai_compatible",
			Models: []organization.OrganizationProviderModel{{
				ID:   "claude-sonnet-4.6",
				Name: "Claude Sonnet",
			}},
		}},
		workspaceOrganization: organization.OrganizationSummary{
			ID:   "org_1",
			Role: role,
		},
		hostedSelection: organization.OrganizationHostedModelSelectionView{
			OrganizationID:         "org_1",
			WorkspaceID:            "wsp_1",
			PickerVisibilityPolicy: pickerPolicy,
			AuthorizationMode:      authorizationMode,
			AuthorizationStatus:    "active",
		},
	}
}

func selectableHostedAvailability(enabled bool, features []string) HostedAvailability {
	return HostedAvailability{
		OrganizationID:    "org_1",
		EntitlementStatus: "eligible_entitled",
		BalanceStatus:     "positive_balance",
		UpdatedAt:         time.Date(2026, 6, 6, 1, 2, 3, 0, time.UTC),
		Models: []HostedAvailabilityModel{{
			HostedModelID:                     "hm_deepseek_v4_pro",
			ModelReference:                    "agentteam/deepseek-v4-pro",
			DisplayName:                       "DeepSeek V4 Pro",
			InventoryStatus:                   "active_inventory",
			EntitlementStatus:                 "eligible_entitled",
			AccessSource:                      "org_grant",
			BalanceStatus:                     "positive_balance",
			RuntimeStatus:                     "runtime_supported",
			Enabled:                           enabled,
			SupportedFamilies:                 []string{"messages", "responses"},
			ExecutableWireAPIFamilies:         []string{"responses"},
			MessagesFeatures:                  append([]string{"tool_use"}, features...),
			ChatCompletionsFeatures:           []string{"json_object"},
			ResponsesFeatures:                 []string{"text_format"},
			InputCreditAtomsPerMillionTokens:  int64Ptr(100),
			OutputCreditAtomsPerMillionTokens: int64Ptr(200),
		}},
	}
}

func selectableHostedPricing() HostedPricing {
	return HostedPricing{
		CreditAtomsPerCredit: 123_456,
		HostedModels: []HostedModelPricing{{
			ModelReference:            "agentteam/deepseek-v4-pro",
			HostedModelID:             "hm_deepseek_v4_pro",
			PricingVersionID:          "hpv_deepseek_v4_pro_discount_90",
			EffectivePricingVersionID: "hpv_deepseek_v4_pro_discount_90",
			PricingFormulaVersion:     "hosted_token_context_threshold_v1",
			BillableDimensions:        []string{"input_tokens", "cache_read_input_tokens", "output_tokens"},
			BasePricingVersionID:      "hpv_deepseek_v4_pro_2026_06_15",
			BasePricingFormulaVersion: "hosted_token_context_threshold_v1",
			BaseBillableDimensions:    []string{"input_tokens", "cache_read_input_tokens", "output_tokens"},
			ContextTiers: []HostedPricingTier{{
				ID:               "short_context",
				MinContextTokens: 0,
				MaxContextTokens: int64Ptr(4096),
				Rates: HostedPricingTierRates{
					InputTokensCreditAtomsPerMillionTokens:    2_000_000,
					CacheReadInputCreditAtomsPerMillionTokens: 200_000,
					OutputCreditAtomsPerMillionTokens:         8_000_000,
				},
			}},
			BaseContextTiers: []HostedPricingTier{{
				ID:               "short_context",
				MinContextTokens: 0,
				MaxContextTokens: int64Ptr(4096),
				Rates: HostedPricingTierRates{
					InputTokensCreditAtomsPerMillionTokens:    4_000_000,
					CacheReadInputCreditAtomsPerMillionTokens: 400_000,
					OutputCreditAtomsPerMillionTokens:         16_000_000,
				},
			}},
			InputCreditAtomsPerMillionTokens:      int64Ptr(2_000_000),
			OutputCreditAtomsPerMillionTokens:     int64Ptr(8_000_000),
			BaseInputCreditAtomsPerMillionTokens:  int64Ptr(4_000_000),
			BaseOutputCreditAtomsPerMillionTokens: int64Ptr(16_000_000),
			PricingOverride: &HostedPricingOverride{
				ID:                        "hpod_deepseek_v4_pro_campaign",
				HostedModelID:             "hm_deepseek_v4_pro",
				ScopeType:                 "global",
				BasePricingVersionID:      "hpv_deepseek_v4_pro_2026_06_15",
				EffectivePricingVersionID: "hpv_deepseek_v4_pro_discount_90",
				OverrideType:              "campaign",
				DisplayLabel:              "90% campaign",
				Priority:                  100,
				Status:                    "active",
			},
		}},
	}
}

func selectableModelRequest(t *testing.T, query string) api.WorkspaceScopedRequest {
	t.Helper()

	recorder := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(recorder)
	ctx.Request = httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_1/selectable-models"+query, nil)
	return api.WorkspaceScopedRequest{
		AuthenticatedRequest: api.AuthenticatedRequest{
			Gin:  ctx,
			User: identity.AuthenticatedUser{User: identity.User{ID: "usr_1"}},
		},
		Workspace: workspace.Summary{WorkspaceID: "wsp_1"},
	}
}

func requireSelectableModel(t *testing.T, models []dto.SelectableModelOption, source string) dto.SelectableModelOption {
	t.Helper()

	for _, model := range models {
		if model.Source == source {
			return model
		}
	}
	t.Fatalf("selectable model source %q not found in %#v", source, models)
	return dto.SelectableModelOption{}
}