base/internal/repository/postgres/asset/asset.go

package asset

import (
	"context"
	"errors"

	"github.com/google/uuid"
	"go.uber.org/fx"
	"gorm.io/gorm"

	domainAsset "base/internal/domain/asset"
)

type assetRepository struct {
	db *gorm.DB
}

func NewAssetRepository(lc fx.Lifecycle, db *gorm.DB) domainAsset.AssetRepository {
	lc.Append(
		fx.Hook{
			OnStart: func(ctx context.Context) error {
				return nil
			},
			OnStop: func(ctx context.Context) error {
				return nil
			},
		})
	return &assetRepository{db: db}
}

func (r *assetRepository) loadRelatedData(ctx context.Context, assetID, categoryID uuid.UUID) (*domainAsset.Category, []domainAsset.Artifact, []domainAsset.Comment, []domainAsset.Report, error) {
	var category *domainAsset.Category
	if categoryID != uuid.Nil {
		var catModel CategoryModel
		if err := r.db.WithContext(ctx).Where("id = ?", categoryID).First(&catModel).Error; err == nil {
			category = toCategoryDomain(&catModel)
		}
	}

	var artifactModels []ArtifactModel
	if err := r.db.WithContext(ctx).Where("asset_id = ?", assetID).Find(&artifactModels).Error; err != nil {
		return nil, nil, nil, nil, err
	}
	artifacts := toArtifactDomains(artifactModels)

	var commentModels []CommentModel
	if err := r.db.WithContext(ctx).Where("asset_id = ?", assetID).Find(&commentModels).Error; err != nil {
		return nil, nil, nil, nil, err
	}
	comments := toCommentDomains(commentModels)

	var reportModels []ReportModel
	if err := r.db.WithContext(ctx).Where("asset_id = ?", assetID).Find(&reportModels).Error; err != nil {
		return nil, nil, nil, nil, err
	}
	reports, err := toReportDomains(reportModels)
	if err != nil {
		return nil, nil, nil, nil, err
	}

	return category, artifacts, comments, reports, nil
}

func (r *assetRepository) Create(ctx context.Context, asset *domainAsset.Asset) error {
	model := toAssetModel(asset)

	tx := r.db.WithContext(ctx).Begin()
	if tx.Error != nil {
		return tx.Error
	}
	defer tx.Rollback()

	if err := tx.Create(model).Error; err != nil {
		return err
	}

	if len(asset.AssetArtifacts) > 0 {
		artifactModels := toArtifactModels(model.ID, asset.AssetArtifacts)
		if err := tx.Create(&artifactModels).Error; err != nil {
			return err
		}
	}

	if len(asset.Comments) > 0 {
		commentModels := toCommentModels(model.ID, asset.Comments)
		if err := tx.Create(&commentModels).Error; err != nil {
			return err
		}
	}

	if len(asset.Reports) > 0 {
		reportModels, err := toReportModels(model.ID, asset.Reports)
		if err != nil {
			return err
		}
		if err := tx.Create(&reportModels).Error; err != nil {
			return err
		}
	}

	if err := tx.Commit().Error; err != nil {
		return err
	}

	copyAssetFromModel(asset, model)
	return nil
}

func (r *assetRepository) FindByID(ctx context.Context, id uuid.UUID) (*domainAsset.Asset, error) {
	var model Model
	if err := r.db.WithContext(ctx).Where("id = ?", id).First(&model).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, errors.New("asset not found")
		}
		return nil, err
	}

	category, artifacts, comments, reports, err := r.loadRelatedData(ctx, model.ID, model.AssetCategoryID)
	if err != nil {
		return nil, err
	}

	return toAssetDomain(&model, category, artifacts, comments, reports), nil
}

func (r *assetRepository) Update(ctx context.Context, asset *domainAsset.Asset) error {
	model := toAssetModel(asset)

	tx := r.db.WithContext(ctx).Begin()
	if tx.Error != nil {
		return tx.Error
	}
	defer tx.Rollback()

	if err := tx.Model(&Model{}).Where("id = ?", asset.ID).Updates(model).Error; err != nil {
		return err
	}

	if err := tx.Where("asset_id = ?", asset.ID).Delete(&ArtifactModel{}).Error; err != nil {
		return err
	}
	if err := tx.Where("asset_id = ?", asset.ID).Delete(&CommentModel{}).Error; err != nil {
		return err
	}
	if err := tx.Where("asset_id = ?", asset.ID).Delete(&ReportModel{}).Error; err != nil {
		return err
	}

	if len(asset.AssetArtifacts) > 0 {
		artifactModels := toArtifactModels(asset.ID, asset.AssetArtifacts)
		if err := tx.Create(&artifactModels).Error; err != nil {
			return err
		}
	}
	if len(asset.Comments) > 0 {
		commentModels := toCommentModels(asset.ID, asset.Comments)
		if err := tx.Create(&commentModels).Error; err != nil {
			return err
		}
	}
	if len(asset.Reports) > 0 {
		reportModels, err := toReportModels(asset.ID, asset.Reports)
		if err != nil {
			return err
		}
		if err := tx.Create(&reportModels).Error; err != nil {
			return err
		}
	}

	return tx.Commit().Error
}

func (r *assetRepository) Delete(ctx context.Context, asset *domainAsset.Asset) error {
	tx := r.db.WithContext(ctx).Begin()
	if tx.Error != nil {
		return tx.Error
	}
	defer tx.Rollback()

	if err := tx.Where("asset_id = ?", asset.ID).Delete(&ArtifactModel{}).Error; err != nil {
		return err
	}
	if err := tx.Where("asset_id = ?", asset.ID).Delete(&CommentModel{}).Error; err != nil {
		return err
	}
	if err := tx.Where("asset_id = ?", asset.ID).Delete(&ReportModel{}).Error; err != nil {
		return err
	}
	if err := tx.Delete(&Model{}, "id = ?", asset.ID).Error; err != nil {
		return err
	}

	return tx.Commit().Error
}

func (r *assetRepository) FindLatestByCategory(ctx context.Context, categoryID uuid.UUID, limit int) ([]*domainAsset.Asset, error) {
	return r.FindLatestByCategoryPaginated(ctx, categoryID, limit, 0)
}

func (r *assetRepository) FindLatestByCategoryPaginated(ctx context.Context, categoryID uuid.UUID, limit, offset int) ([]*domainAsset.Asset, error) {
	var models []Model
	q := r.db.WithContext(ctx).Where("asset_category_id = ?", categoryID).Order("created_at DESC")

	if limit > 0 {
		q = q.Limit(limit)
	}

	if offset > 0 {
		q = q.Offset(offset)
	}

	if err := q.Find(&models).Error; err != nil {
		return nil, err
	}

	if len(models) == 0 {
		return nil, nil
	}

	out := make([]*domainAsset.Asset, len(models))

	for i, model := range models {
		category, artifacts, comments, reports, err := r.loadRelatedData(ctx, model.ID, model.AssetCategoryID)
		if err != nil {
			return nil, err
		}
		out[i] = toAssetDomain(&model, category, artifacts, comments, reports)
	}

	return out, nil
}

func (r *assetRepository) CountByCategory(ctx context.Context, categoryID uuid.UUID) (int, error) {
	var count int64
	if err := r.db.WithContext(ctx).Model(&Model{}).Where("asset_category_id = ?", categoryID).Count(&count).Error; err != nil {
		return 0, err
	}
	return int(count), nil
}

func (r *assetRepository) FindLatest(ctx context.Context, limit, offset int) ([]*domainAsset.Asset, error) {
	var models []Model
	q := r.db.WithContext(ctx).Order("created_at DESC")
	if limit > 0 {
		q = q.Limit(limit)
	}
	if offset > 0 {
		q = q.Offset(offset)
	}
	if err := q.Find(&models).Error; err != nil {
		return nil, err
	}

	if len(models) == 0 {
		return nil, nil
	}

	out := make([]*domainAsset.Asset, len(models))
	for i, model := range models {
		category, artifacts, comments, reports, err := r.loadRelatedData(ctx, model.ID, model.AssetCategoryID)
		if err != nil {
			return nil, err
		}
		out[i] = toAssetDomain(&model, category, artifacts, comments, reports)
	}
	return out, nil
}

func (r *assetRepository) FindByProfileID(ctx context.Context, profileID uuid.UUID) ([]*domainAsset.Asset, error) {
	var models []Model
	if err := r.db.WithContext(ctx).Where("profile_id = ?", profileID).Order("created_at DESC").Find(&models).Error; err != nil {
		return nil, err
	}

	if len(models) == 0 {
		return nil, nil
	}

	out := make([]*domainAsset.Asset, len(models))
	for i, model := range models {
		category, artifacts, comments, reports, err := r.loadRelatedData(ctx, model.ID, model.AssetCategoryID)
		if err != nil {
			return nil, err
		}
		out[i] = toAssetDomain(&model, category, artifacts, comments, reports)
	}
	return out, nil
}

func (r *assetRepository) Count(ctx context.Context) (int, error) {
	var count int64
	if err := r.db.WithContext(ctx).Model(&Model{}).Count(&count).Error; err != nil {
		return 0, err
	}
	return int(count), nil
}