package repository

import (
	"context"

	"gorm.io/gorm"

	vulnmodel "security-center-backend/internal/modules/vuln/model"
)

type SLARuleRepository interface {
	List(ctx context.Context) ([]vulnmodel.SLARule, error)
	Upsert(ctx context.Context, rules map[string]int) error
}

type GormSLARuleRepository struct {
	db *gorm.DB
}

func NewSLARuleRepository(db *gorm.DB) SLARuleRepository {
	return &GormSLARuleRepository{db: db}
}

func (r *GormSLARuleRepository) List(ctx context.Context) ([]vulnmodel.SLARule, error) {
	var rules []vulnmodel.SLARule
	if err := r.db.WithContext(ctx).Find(&rules).Error; err != nil {
		return nil, err
	}
	return rules, nil
}

func (r *GormSLARuleRepository) Upsert(ctx context.Context, rules map[string]int) error {
	for severity, hours := range rules {
		rule := vulnmodel.SLARule{Severity: severity, Hours: hours}
		if err := r.db.WithContext(ctx).Where("severity = ?", severity).Assign(rule).FirstOrCreate(&rule).Error; err != nil {
			return err
		}
	}
	return nil
}