package server

import (
	"context"
	"fmt"

	"github.com/tidbcloud/AgentTeam/internal/api/billing"
	"github.com/tidbcloud/AgentTeam/internal/api/workspace"
)

type workspaceHostedPricingClient struct {
	next interface {
		LoadHostedPricing(context.Context, string) (billing.HostedPricing, error)
	}
}

func (client workspaceHostedPricingClient) LoadHostedPricing(ctx context.Context, organizationID string) (workspace.HostedPricing, error) {
	if client.next == nil {
		return workspace.HostedPricing{}, fmt.Errorf("agentteam-maas hosted pricing client is not configured")
	}
	pricing, err := client.next.LoadHostedPricing(ctx, organizationID)
	if err != nil {
		return workspace.HostedPricing{}, err
	}
	models := make([]workspace.HostedModelPricing, 0, len(pricing.HostedModels))
	for _, model := range pricing.HostedModels {
		models = append(models, workspace.HostedModelPricing{
			ModelReference:                                model.ModelReference,
			HostedModelID:                                 model.HostedModelID,
			PricingVersionID:                              model.PricingVersionID,
			EffectivePricingVersionID:                     model.EffectivePricingVersionID,
			PricingFormulaVersion:                         model.PricingFormulaVersion,
			BillableDimensions:                            append([]string(nil), model.BillableDimensions...),
			ContextTiers:                                  workspaceHostedPricingTiers(model.ContextTiers),
			BasePricingVersionID:                          model.BasePricingVersionID,
			BasePricingFormulaVersion:                     model.BasePricingFormulaVersion,
			BaseBillableDimensions:                        append([]string(nil), model.BaseBillableDimensions...),
			BaseContextTiers:                              workspaceHostedPricingTiers(model.BaseContextTiers),
			InputCreditAtomsPerMillionTokens:              model.InputCreditAtomsPerMillionTokens,
			CacheReadInputCreditAtomsPerMillionTokens:     model.CacheReadInputCreditAtomsPerMillionTokens,
			OutputCreditAtomsPerMillionTokens:             model.OutputCreditAtomsPerMillionTokens,
			BaseInputCreditAtomsPerMillionTokens:          model.BaseInputCreditAtomsPerMillionTokens,
			BaseCacheReadInputCreditAtomsPerMillionTokens: model.BaseCacheReadInputCreditAtomsPerMillionTokens,
			BaseOutputCreditAtomsPerMillionTokens:         model.BaseOutputCreditAtomsPerMillionTokens,
			PricingOverride:                               workspaceHostedPricingOverride(model.PricingOverride),
		})
	}
	return workspace.HostedPricing{
		CreditAtomsPerCredit: pricing.CreditAtomsPerCredit,
		HostedModels:         models,
	}, nil
}

func workspaceHostedPricingTiers(tiers []billing.HostedPricingTier) []workspace.HostedPricingTier {
	if len(tiers) == 0 {
		return nil
	}
	copied := make([]workspace.HostedPricingTier, 0, len(tiers))
	for _, tier := range tiers {
		copied = append(copied, workspace.HostedPricingTier{
			ID:               tier.ID,
			MinContextTokens: tier.MinContextTokens,
			MaxContextTokens: tier.MaxContextTokens,
			Rates: workspace.HostedPricingTierRates{
				InputTokensCreditAtomsPerMillionTokens:    tier.Rates.InputTokensCreditAtomsPerMillionTokens,
				CacheReadInputCreditAtomsPerMillionTokens: tier.Rates.CacheReadInputCreditAtomsPerMillionTokens,
				OutputCreditAtomsPerMillionTokens:         tier.Rates.OutputCreditAtomsPerMillionTokens,
			},
		})
	}
	return copied
}

func workspaceHostedPricingOverride(override *billing.HostedPricingOverride) *workspace.HostedPricingOverride {
	if override == nil {
		return nil
	}
	return &workspace.HostedPricingOverride{
		ID:                        override.ID,
		HostedModelID:             override.HostedModelID,
		ScopeType:                 override.ScopeType,
		ScopeID:                   override.ScopeID,
		BasePricingVersionID:      override.BasePricingVersionID,
		EffectivePricingVersionID: override.EffectivePricingVersionID,
		OverrideType:              override.OverrideType,
		DisplayLabel:              override.DisplayLabel,
		Priority:                  override.Priority,
		EffectiveFrom:             override.EffectiveFrom,
		EffectiveTo:               override.EffectiveTo,
		Status:                    override.Status,
	}
}