package repository

import (
	"context"

	"gorm.io/gorm"

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

type VulnerabilityAttachmentRepository interface {
	Create(ctx context.Context, attachment *vulnmodel.VulnerabilityAttachment) error
	FindByID(ctx context.Context, id uint64) (*vulnmodel.VulnerabilityAttachment, error)
	UpdateFields(ctx context.Context, id uint64, fields map[string]interface{}) error
	AttachToVulnerability(ctx context.Context, uploaderID uint64, vulnID uint64, attachmentIDs []uint64) error
}

type GormVulnerabilityAttachmentRepository struct {
	db *gorm.DB
}

func NewVulnerabilityAttachmentRepository(db *gorm.DB) VulnerabilityAttachmentRepository {
	return &GormVulnerabilityAttachmentRepository{db: db}
}

func (r *GormVulnerabilityAttachmentRepository) Create(ctx context.Context, attachment *vulnmodel.VulnerabilityAttachment) error {
	return r.db.WithContext(ctx).Create(attachment).Error
}

func (r *GormVulnerabilityAttachmentRepository) FindByID(ctx context.Context, id uint64) (*vulnmodel.VulnerabilityAttachment, error) {
	var item vulnmodel.VulnerabilityAttachment
	if err := r.db.WithContext(ctx).First(&item, id).Error; err != nil {
		return nil, err
	}
	return &item, nil
}

func (r *GormVulnerabilityAttachmentRepository) UpdateFields(ctx context.Context, id uint64, fields map[string]interface{}) error {
	return r.db.WithContext(ctx).Model(&vulnmodel.VulnerabilityAttachment{}).Where("id = ?", id).Updates(fields).Error
}

func (r *GormVulnerabilityAttachmentRepository) AttachToVulnerability(ctx context.Context, uploaderID uint64, vulnID uint64, attachmentIDs []uint64) error {
	if len(attachmentIDs) == 0 {
		return nil
	}
	return r.db.WithContext(ctx).
		Model(&vulnmodel.VulnerabilityAttachment{}).
		Where(
			"id IN ? AND uploader_id = ? AND status = ? AND (vulnerability_id IS NULL OR vulnerability_id = 0)",
			attachmentIDs,
			uploaderID,
			vulnmodel.AttachmentStatusReady,
		).
		Updates(map[string]interface{}{
			"vulnerability_id": vulnID,
		}).Error
}