package ai

import (
	"context"
	"strings"

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

	"github.com/ultisuite/ulti-backend/internal/ai/cost"
	"github.com/ultisuite/ulti-backend/internal/llm"
)

// QuotaService wraps cost policy and metering for backward compatibility.
type QuotaService struct {
	db     *pgxpool.Pool
	policy *cost.PolicyService
	meter  *cost.Meter
}

func NewQuotaService(db *pgxpool.Pool) *QuotaService {
	return &QuotaService{
		db:     db,
		policy: cost.NewPolicyService(db),
		meter:  cost.NewMeter(db),
	}
}

func (s *QuotaService) Check(ctx context.Context, externalUserID string) (SpendStatus, error) {
	orgBilling, _ := s.usesOrgBilling(ctx, externalUserID)
	status, err := s.policy.GetStatus(ctx, externalUserID, orgBilling)
	if err != nil {
		return SpendStatus{}, err
	}
	return status, nil
}

func (s *QuotaService) AssertAvailable(ctx context.Context, externalUserID string, provider llm.Provider, useOrgSettings bool) error {
	scope := ResolveBillingScope(ctx, s.db, externalUserID, provider, useOrgSettings)
	return s.policy.AssertAvailable(ctx, externalUserID, scope)
}

func (s *QuotaService) RecordUsage(ctx context.Context, in cost.RecordInput) error {
	return s.meter.RecordUsage(ctx, in)
}

func (s *QuotaService) usesOrgBilling(ctx context.Context, externalUserID string) (bool, error) {
	settings, err := LoadEffectiveLLMSettings(ctx, s.db, externalUserID)
	if err != nil {
		return true, err
	}
	org, err := loadOrgLLMPolicy(ctx, s.db)
	if err != nil {
		return true, err
	}
	if org.EnforceOrgProviders || len(settings.Providers) == 0 {
		return true, nil
	}
	user, err := loadUserLLMSettings(ctx, s.db, externalUserID)
	if err != nil {
		return true, err
	}
	return len(user.Providers) == 0, nil
}

// ResolveBillingScope determines whether usage is billed to org or user's own key.
func ResolveBillingScope(ctx context.Context, db *pgxpool.Pool, externalUserID string, provider llm.Provider, useOrgSettings bool) string {
	if useOrgSettings {
		return cost.BillingScopeOrg
	}
	org, err := loadOrgLLMPolicy(ctx, db)
	if err != nil || org.EnforceOrgProviders {
		return cost.BillingScopeOrg
	}
	user, err := loadUserLLMSettings(ctx, db, externalUserID)
	if err != nil || len(user.Providers) == 0 {
		return cost.BillingScopeOrg
	}
	apiKey := strings.TrimSpace(provider.APIKey)
	for _, op := range org.Providers {
		if op.ID == provider.ID && strings.TrimSpace(op.APIKey) == apiKey && apiKey != "" {
			return cost.BillingScopeOrg
		}
	}
	for _, up := range user.Providers {
		if up.ID == provider.ID && strings.TrimSpace(up.APIKey) != "" {
			return cost.BillingScopeUser
		}
	}
	return cost.BillingScopeOrg
}

// SpendStatus is the user-facing quota/spend response.
type SpendStatus = cost.SpendStatus

// ErrQuotaExceeded aliases cost limit error for backward compatibility.
var ErrQuotaExceeded = cost.ErrCostLimitExceeded