package quota

import (
	"errors"
	"log/slog"
	"time"

	"github.com/google/uuid"
	"gorm.io/gorm"

	"github.com/llamalink/llamalink/internal/db"
)

var (
	ErrQuotaExceeded = errors.New("quota exceeded for this period")
)

type Service struct {
	db *gorm.DB
}

func NewService(db *gorm.DB) *Service {
	return &Service{db: db}
}

func (s *Service) CheckQuota(apiKeyID uuid.UUID, modelScope string) (bool, string, error) {
	now := time.Now().UTC()

	var quota db.Quota
	err := s.db.Where(
		"api_key_id = ? AND period_start <= ? AND period_end > ?",
		apiKeyID, now, now,
	).First(&quota).Error

	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return true, "", nil // No quota configured
		}
		return false, "", err
	}

	// Check model-specific scope if set
	if quota.ModelScope != nil && *quota.ModelScope != "" && modelScope != "" {
		if *quota.ModelScope != modelScope {
			return true, "", nil // Quota doesn't apply to this model
		}
	}

	if quota.TokensUsed >= quota.TokensLimit {
		return false, "quota exceeded", nil
	}

	// Warning at 90%
	if quota.TokensLimit > 0 {
		usage := float64(quota.TokensUsed) / float64(quota.TokensLimit)
		if usage >= 0.9 {
			slog.Warn("quota usage warning",
				"api_key_id", apiKeyID,
				"usage_percent", usage*100,
				"tokens_used", quota.TokensUsed,
				"tokens_limit", quota.TokensLimit,
			)
		}
	}

	return true, "", nil
}

func (s *Service) ConsumeQuota(apiKeyID uuid.UUID, modelScope string, tokens int) error {
	now := time.Now().UTC()
	periodStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC)
	periodEnd := periodStart.AddDate(0, 1, 0)

	return s.db.Transaction(func(tx *gorm.DB) error {
		var quota db.Quota
		err := tx.Where(
			"api_key_id = ? AND period_start = ? AND period_end = ?",
			apiKeyID, periodStart, periodEnd,
		).First(&quota).Error

		if err != nil {
			if errors.Is(err, gorm.ErrRecordNotFound) {
				return nil // No quota configured
			}
			return err
		}

		// Check model scope
		if quota.ModelScope != nil && *quota.ModelScope != "" && modelScope != "" {
			if *quota.ModelScope != modelScope {
				return nil // Quota doesn't apply
			}
		}

		newUsed := quota.TokensUsed + tokens
		if quota.TokensLimit > 0 && newUsed > quota.TokensLimit {
			return ErrQuotaExceeded
		}

		return tx.Model(&quota).Update("tokens_used", newUsed).Error
	})
}

func (s *Service) GetUsage(apiKeyID uuid.UUID, periodStart, periodEnd time.Time) (used int, limit int, err error) {
	var quota db.Quota
	err = s.db.Where(
		"api_key_id = ? AND period_start = ? AND period_end = ?",
		apiKeyID, periodStart, periodEnd,
	).First(&quota).Error

	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return 0, 0, nil
		}
		return 0, 0, err
	}

	return quota.TokensUsed, quota.TokensLimit, nil
}