package store

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"strings"

	"github.com/jackc/pgx/v5"
)

type platformModelQuerier interface {
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

type modelCatalogSnapshot struct {
	ID                     string
	ProviderKey            string
	CanonicalModelKey      string
	InvocationName         string
	ProviderModelName      string
	ModelType              StringList
	DisplayName            string
	Capabilities           map[string]any
	BaseBillingConfig      map[string]any
	PricingRuleSetID       string
	DefaultRateLimitPolicy map[string]any
	RuntimePolicySetID     string
	RuntimePolicyOverride  map[string]any
	Status                 string
}

func (s *Store) CreatePlatformModel(ctx context.Context, input CreatePlatformModelInput) (PlatformModel, error) {
	return s.createPlatformModel(ctx, s.pool, input)
}

// EffectivePlatformModelCapabilities merges base defaults with platform overrides.
// Nested capability objects are merged recursively so platform fields override
// only the same nested keys while preserving other base capability defaults.
func EffectivePlatformModelCapabilities(base map[string]any, override map[string]any) map[string]any {
	return mergeCapabilityObjects(base, override)
}

func (s *Store) ReplacePlatformModels(ctx context.Context, platformID string, inputs []CreatePlatformModelInput) ([]PlatformModel, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer rollbackTransaction(tx)

	keptIDs := make([]string, 0, len(inputs))
	for _, input := range inputs {
		input.PlatformID = platformID
		model, err := s.createPlatformModel(ctx, tx, input)
		if err != nil {
			return nil, err
		}
		keptIDs = append(keptIDs, model.ID)
	}

	if len(keptIDs) == 0 {
		if _, err := tx.Exec(ctx, `
WITH deleted AS (
  DELETE FROM platform_models
  WHERE platform_id = $1::uuid
  RETURNING id
)
DELETE FROM gateway_access_rules
WHERE resource_type = 'platform_model'
  AND resource_id IN (SELECT id FROM deleted)`, platformID); err != nil {
			return nil, err
		}
	} else if _, err := tx.Exec(ctx, `
WITH deleted AS (
  DELETE FROM platform_models
  WHERE platform_id = $1::uuid
    AND NOT (id::text = ANY($2::text[]))
  RETURNING id
)
DELETE FROM gateway_access_rules
WHERE resource_type = 'platform_model'
  AND resource_id IN (SELECT id FROM deleted)`, platformID, keptIDs); err != nil {
		return nil, err
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, err
	}
	return s.ListPlatformModels(ctx, platformID)
}

func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier, input CreatePlatformModelInput) (PlatformModel, error) {
	input.ModelName = strings.TrimSpace(input.ModelName)
	input.ProviderModelName = strings.TrimSpace(input.ProviderModelName)
	normalizedRateLimitPolicy, err := NormalizeAndValidateRateLimitPolicy(input.RateLimitPolicy)
	if err != nil {
		return PlatformModel{}, err
	}
	input.RateLimitPolicy = normalizedRateLimitPolicy
	if rawPolicy, ok := input.RuntimePolicyOverride["rateLimitPolicy"].(map[string]any); ok {
		normalizedRuntimeOverridePolicy, normalizeErr := NormalizeAndValidateRateLimitPolicy(rawPolicy)
		if normalizeErr != nil {
			return PlatformModel{}, normalizeErr
		}
		input.RuntimePolicyOverride = cloneObject(input.RuntimePolicyOverride)
		input.RuntimePolicyOverride["rateLimitPolicy"] = normalizedRuntimeOverridePolicy
	}
	base, err := s.lookupBaseModel(ctx, q, input.BaseModelID, input.CanonicalModelKey, input.ModelName)
	if err != nil && !IsNotFound(err) {
		return PlatformModel{}, err
	}
	if base.ID != "" && base.Status != "" && base.Status != "active" {
		var existingBinding bool
		if err := q.QueryRow(ctx, `
SELECT EXISTS (
  SELECT 1 FROM platform_models WHERE platform_id = $1::uuid AND base_model_id = $2::uuid
)`, input.PlatformID, base.ID).Scan(&existingBinding); err != nil {
			return PlatformModel{}, err
		}
		if !existingBinding {
			return PlatformModel{}, fmt.Errorf("%w: base model %q is %s and cannot be newly bound", ErrInvalidPlatformModelConfiguration, base.InvocationName, base.Status)
		}
	}
	if len(input.ModelType) == 0 {
		input.ModelType = base.ModelType
	}
	input.ModelType = normalizeModelTypeList(input.ModelType)
	if len(input.ModelType) == 0 {
		input.ModelType = StringList{"text_generate"}
	}
	if base.InvocationName != "" {
		input.ModelName = base.InvocationName
	} else if input.ModelName == "" {
		input.ModelName = base.ProviderModelName
	}
	if input.ProviderModelName == "" {
		input.ProviderModelName = input.ModelName
	}
	if input.DisplayName == "" {
		input.DisplayName = firstNonEmpty(base.DisplayName, input.ModelName)
	}
	input.ModelAlias = normalizePlatformModelAlias(input.ModelAlias, base)
	if input.PricingMode == "" {
		input.PricingMode = "inherit_discount"
	}
	capabilities := input.Capabilities
	if len(capabilities) == 0 {
		capabilities = EffectivePlatformModelCapabilities(base.Capabilities, input.CapabilityOverride)
	}
	if err := validateEnabledVolcesTextModelCapabilities(ctx, q, input, capabilities); err != nil {
		return PlatformModel{}, err
	}
	// billing_config is a legacy, explicitly supplied compatibility field. Do
	// not materialize base-model pricing into it: copied prices become stale as
	// soon as the base pricing rule changes and can mask the authoritative rule.
	billingConfig := input.BillingConfig
	explicitRuntimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID)
	rateLimitPolicyMode := NormalizeRateLimitPolicyMode(input.RateLimitPolicyMode, input.RateLimitPolicy != nil)
	rateLimitPolicy := input.RateLimitPolicy
	if rateLimitPolicyMode == RateLimitPolicyModeInherit {
		rateLimitPolicy = nil
	}
	runtimePolicySetID := explicitRuntimePolicySetID
	runtimePolicyOverride := input.RuntimePolicyOverride

	capabilityOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.CapabilityOverride))
	capabilitiesJSON, _ := json.Marshal(emptyObjectIfNil(capabilities))
	modelTypeJSON, _ := json.Marshal(input.ModelType)
	billingOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingConfigOverride))
	billingJSON, _ := json.Marshal(emptyObjectIfNil(billingConfig))
	permissionJSON, _ := json.Marshal(emptyObjectIfNil(input.PermissionConfig))
	retryJSON, _ := json.Marshal(emptyObjectIfNil(input.RetryPolicy))
	rateLimitJSON, _ := json.Marshal(emptyObjectIfNil(rateLimitPolicy))
	runtimePolicyOverrideJSON, _ := json.Marshal(emptyObjectIfNil(runtimePolicyOverride))

	discount := any(nil)
	if input.DiscountFactor > 0 {
		discount = input.DiscountFactor
	}
	baseID := any(nil)
	if base.ID != "" {
		baseID = base.ID
	}

	var model PlatformModel
	var capabilityOverrideBytes []byte
	var capabilitiesBytes []byte
	var billingOverrideBytes []byte
	var billingBytes []byte
	var permissionBytes []byte
	var retryPolicyBytes []byte
	var rateLimitPolicyBytes []byte
	var runtimePolicyOverrideBytes []byte
	var modelTypeBytes []byte
	err = q.QueryRow(ctx, `
INSERT INTO platform_models (
  platform_id, base_model_id, model_name, provider_model_name, model_alias, model_type, display_name,
  capability_override, capabilities, pricing_mode, discount_factor,
  pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode,
  runtime_policy_set_id, runtime_policy_override, enabled
)
VALUES (
  $1::uuid, $2::uuid, $3, NULLIF($4, ''), NULLIF($5, ''), $6::jsonb, $7,
  $8::jsonb, $9::jsonb, $10, $11::numeric,
  NULLIF($12, '')::uuid, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, $17::jsonb, $18,
  NULLIF($19, '')::uuid, $20::jsonb, true
)
ON CONFLICT (platform_id, model_name) DO UPDATE
SET base_model_id = EXCLUDED.base_model_id,
    provider_model_name = EXCLUDED.provider_model_name,
    model_alias = EXCLUDED.model_alias,
    display_name = EXCLUDED.display_name,
    capability_override = EXCLUDED.capability_override,
    capabilities = EXCLUDED.capabilities,
    pricing_mode = EXCLUDED.pricing_mode,
    discount_factor = EXCLUDED.discount_factor,
    pricing_rule_set_id = EXCLUDED.pricing_rule_set_id,
    billing_config_override = EXCLUDED.billing_config_override,
    billing_config = EXCLUDED.billing_config,
    permission_config = EXCLUDED.permission_config,
    retry_policy = EXCLUDED.retry_policy,
    rate_limit_policy = EXCLUDED.rate_limit_policy,
    rate_limit_policy_mode = EXCLUDED.rate_limit_policy_mode,
    runtime_policy_set_id = EXCLUDED.runtime_policy_set_id,
    runtime_policy_override = EXCLUDED.runtime_policy_override,
    enabled = true,
    updated_at = now()
RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_name,
          COALESCE(NULLIF(provider_model_name, ''), model_name), COALESCE(model_alias, ''), model_type, display_name, capability_override,
          capabilities, pricing_mode, COALESCE(discount_factor, 0)::float8,
          COALESCE(pricing_rule_set_id::text, ''), billing_config_override, billing_config,
          permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override,
          COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
          enabled, created_at, updated_at`,
		input.PlatformID,
		baseID,
		input.ModelName,
		input.ProviderModelName,
		input.ModelAlias,
		string(modelTypeJSON),
		input.DisplayName,
		string(capabilityOverrideJSON),
		string(capabilitiesJSON),
		input.PricingMode,
		discount,
		input.PricingRuleSetID,
		string(billingOverrideJSON),
		string(billingJSON),
		string(permissionJSON),
		string(retryJSON),
		string(rateLimitJSON),
		rateLimitPolicyMode,
		runtimePolicySetID,
		string(runtimePolicyOverrideJSON),
	).Scan(
		&model.ID,
		&model.PlatformID,
		&model.BaseModelID,
		&model.ModelName,
		&model.ProviderModelName,
		&model.ModelAlias,
		&modelTypeBytes,
		&model.DisplayName,
		&capabilityOverrideBytes,
		&capabilitiesBytes,
		&model.PricingMode,
		&model.DiscountFactor,
		&model.PricingRuleSetID,
		&billingOverrideBytes,
		&billingBytes,
		&permissionBytes,
		&retryPolicyBytes,
		&rateLimitPolicyBytes,
		&model.RateLimitPolicyMode,
		&model.RuntimePolicySetID,
		&runtimePolicyOverrideBytes,
		&model.CooldownUntil,
		&model.Enabled,
		&model.CreatedAt,
		&model.UpdatedAt,
	)
	if err != nil {
		return PlatformModel{}, err
	}
	model.CapabilityOverride = decodeObject(capabilityOverrideBytes)
	model.Capabilities = decodeObject(capabilitiesBytes)
	model.ModelType = decodeStringArray(modelTypeBytes)
	model.BillingConfigOverride = decodeObject(billingOverrideBytes)
	model.BillingConfig = decodeObject(billingBytes)
	model.BaseBillingConfig = base.BaseBillingConfig
	model.BasePricingRuleSetID = base.PricingRuleSetID
	model.PermissionConfig = decodeObject(permissionBytes)
	model.RetryPolicy = decodeObject(retryPolicyBytes)
	model.RateLimitPolicy = decodeObject(rateLimitPolicyBytes)
	model.RuntimePolicyOverride = decodeObject(runtimePolicyOverrideBytes)
	return model, nil
}

func validateEnabledVolcesTextModelCapabilities(ctx context.Context, q platformModelQuerier, input CreatePlatformModelInput, capabilities map[string]any) error {
	// createPlatformModel enables/upserts models unconditionally, so every text
	// model that reaches this path must already satisfy the Volcengine invariant.
	if !containsTextOutputModelType(input.ModelType) {
		return nil
	}
	var provider string
	var baseURL string
	if err := q.QueryRow(ctx, `SELECT provider, COALESCE(base_url, '') FROM integration_platforms WHERE id = $1::uuid`, input.PlatformID).Scan(&provider, &baseURL); err != nil {
		return err
	}
	provider = strings.ToLower(strings.TrimSpace(provider))
	baseURL = strings.ToLower(strings.TrimSpace(baseURL))
	if provider != "volces-openai" && !strings.Contains(baseURL, "volces.com") && !strings.Contains(baseURL, "byteplus.com") {
		return nil
	}
	seen := map[string]struct{}{}
	for _, modelType := range append(append(StringList{}, input.ModelType...), "text_generate") {
		modelType = strings.TrimSpace(modelType)
		if modelType == "" {
			continue
		}
		if _, ok := seen[modelType]; ok {
			continue
		}
		seen[modelType] = struct{}{}
		capability, _ := capabilities[modelType].(map[string]any)
		if value, ok := positiveWholeNumber(capability["max_output_tokens"]); ok && value > 0 {
			return nil
		}
	}
	return fmt.Errorf("%w: enabled Volcengine text model %q requires a positive integer max_output_tokens capability for its text model type or text_generate fallback", ErrInvalidPlatformModelConfiguration, input.ProviderModelName)
}

func containsTextOutputModelType(values StringList) bool {
	for _, value := range values {
		switch strings.ToLower(strings.TrimSpace(value)) {
		case "text_generate", "chat", "responses", "text":
			return true
		}
	}
	return false
}

func positiveWholeNumber(value any) (int64, bool) {
	var number float64
	switch typed := value.(type) {
	case int:
		number = float64(typed)
	case int32:
		number = float64(typed)
	case int64:
		number = float64(typed)
	case float64:
		number = typed
	case json.Number:
		parsed, err := typed.Float64()
		if err != nil {
			return 0, false
		}
		number = parsed
	default:
		return 0, false
	}
	if number <= 0 || math.Trunc(number) != number || number > math.MaxInt64 {
		return 0, false
	}
	return int64(number), true
}

func (s *Store) DeletePlatformModel(ctx context.Context, id string) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollbackTransaction(tx)

	result, err := tx.Exec(ctx, `DELETE FROM platform_models WHERE id = $1::uuid`, id)
	if err != nil {
		return err
	}
	if result.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	if _, err := tx.Exec(ctx, `
DELETE FROM gateway_access_rules
WHERE resource_type = 'platform_model' AND resource_id = $1::uuid`, id); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (s *Store) lookupBaseModel(ctx context.Context, q platformModelQuerier, id string, canonicalKey string, modelName string) (modelCatalogSnapshot, error) {
	var item modelCatalogSnapshot
	var capabilities []byte
	var billingConfig []byte
	var rateLimitPolicy []byte
	var runtimePolicyOverride []byte
	var modelTypeBytes []byte
	err := q.QueryRow(ctx, `
SELECT id::text, provider_key, canonical_model_key, invocation_name, provider_model_name, model_type, display_name,
       capabilities, base_billing_config, COALESCE(pricing_rule_set_id::text, ''), default_rate_limit_policy,
       COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override, status
FROM base_model_catalog
WHERE ($1 <> '' AND id = NULLIF($1, '')::uuid)
   OR ($2 <> '' AND canonical_model_key = $2)
   OR ($3 <> '' AND provider_model_name = $3)
ORDER BY CASE WHEN id::text = $1 THEN 0 WHEN canonical_model_key = $2 THEN 1 ELSE 2 END
LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSpace(modelName)).Scan(
		&item.ID,
		&item.ProviderKey,
		&item.CanonicalModelKey,
		&item.InvocationName,
		&item.ProviderModelName,
		&modelTypeBytes,
		&item.DisplayName,
		&capabilities,
		&billingConfig,
		&item.PricingRuleSetID,
		&rateLimitPolicy,
		&item.RuntimePolicySetID,
		&runtimePolicyOverride,
		&item.Status,
	)
	if err != nil {
		if err == pgx.ErrNoRows {
			return modelCatalogSnapshot{}, err
		}
		return modelCatalogSnapshot{}, err
	}
	item.Capabilities = decodeObject(capabilities)
	item.BaseBillingConfig = decodeObject(billingConfig)
	item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
	item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
	item.ModelType = normalizeModelTypeList(decodeStringArray(modelTypeBytes))
	return item, nil
}

func normalizePlatformModelAlias(alias string, base modelCatalogSnapshot) string {
	if base.InvocationName != "" {
		return base.InvocationName
	}
	return strings.TrimSpace(alias)
}

func stripAliasPrefix(alias string) string {
	if before, after, ok := strings.Cut(alias, ":"); ok && before != "" && after != "" {
		return after
	}
	return alias
}

func mergeObjects(base map[string]any, override map[string]any) map[string]any {
	out := map[string]any{}
	for key, value := range base {
		out[key] = value
	}
	for key, value := range override {
		out[key] = value
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

func mergeCapabilityObjects(base map[string]any, override map[string]any) map[string]any {
	out := cloneObject(base)
	for key, value := range override {
		baseChild, baseOK := out[key].(map[string]any)
		overrideChild, overrideOK := value.(map[string]any)
		if baseOK && overrideOK {
			out[key] = mergeCapabilityObjects(baseChild, overrideChild)
			continue
		}
		out[key] = cloneObjectValue(value)
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

func cloneObject(values map[string]any) map[string]any {
	out := make(map[string]any, len(values))
	for key, value := range values {
		out[key] = cloneObjectValue(value)
	}
	return out
}

func cloneObjectValue(value any) any {
	switch typed := value.(type) {
	case map[string]any:
		return cloneObject(typed)
	case []any:
		out := make([]any, 0, len(typed))
		for _, item := range typed {
			out = append(out, cloneObjectValue(item))
		}
		return out
	case []string:
		return append([]string(nil), typed...)
	default:
		return value
	}
}

func emptyObjectIfNil(value map[string]any) map[string]any {
	if value == nil {
		return map[string]any{}
	}
	return value
}
