package repository import ( "context" "strings" "haixun-backend/internal/library/clock" app "haixun-backend/internal/library/errors" "haixun-backend/internal/library/errors/code" libkg "haixun-backend/internal/library/knowledge" libmongo "haixun-backend/internal/library/mongo" "haixun-backend/internal/model/knowledge_graph/domain/entity" domrepo "haixun-backend/internal/model/knowledge_graph/domain/repository" "github.com/google/uuid" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" ) type mongoRepository struct { collection *mongo.Collection } func NewMongoRepository(db *mongo.Database) domrepo.Repository { if db == nil { return &mongoRepository{} } return &mongoRepository{collection: db.Collection(entity.CollectionName)} } func (r *mongoRepository) EnsureIndexes(ctx context.Context) error { if r.collection == nil { return nil } return libmongo.EnsureIndexes(ctx, r.collection, []mongo.IndexModel{ {Keys: bson.D{{Key: "tenant_id", Value: 1}, {Key: "owner_uid", Value: 1}, {Key: "brand_id", Value: 1}}, Options: options.Index().SetUnique(true).SetPartialFilterExpression(bson.M{"brand_id": bson.M{"$gt": ""}, "topic_id": bson.M{"$in": []interface{}{nil, ""}}})}, {Keys: bson.D{{Key: "tenant_id", Value: 1}, {Key: "owner_uid", Value: 1}, {Key: "topic_id", Value: 1}}, Options: options.Index().SetUnique(true).SetPartialFilterExpression(bson.M{"topic_id": bson.M{"$gt": ""}})}, {Keys: bson.D{{Key: "tenant_id", Value: 1}, {Key: "owner_uid", Value: 1}, {Key: "persona_id", Value: 1}}, Options: options.Index().SetUnique(true).SetPartialFilterExpression(bson.M{"persona_id": bson.M{"$gt": ""}})}, {Keys: bson.D{{Key: "brand_id", Value: 1}, {Key: "update_at", Value: -1}}}, {Keys: bson.D{{Key: "topic_id", Value: 1}, {Key: "update_at", Value: -1}}}, {Keys: bson.D{{Key: "persona_id", Value: 1}, {Key: "update_at", Value: -1}}}, }) } func brandOwnerFilter(tenantID, ownerUID, brandID string) bson.M { filter := bson.M{ "tenant_id": tenantID, "owner_uid": ownerUID, } for k, v := range libmongo.BrandScopeFilter(brandID) { filter[k] = v } return filter } func topicOwnerFilter(tenantID, ownerUID, topicID string) bson.M { return bson.M{ "tenant_id": tenantID, "owner_uid": ownerUID, "topic_id": strings.TrimSpace(topicID), } } func (r *mongoRepository) UpsertByBrand(ctx context.Context, graph *entity.Graph) (*entity.Graph, error) { if r.collection == nil { return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured") } if graph == nil { return nil, app.For(code.Brand).InputMissingRequired("graph is required") } now := clock.NowUnixNano() graph.UpdateAt = now if graph.CreateAt == 0 { graph.CreateAt = now } if strings.TrimSpace(graph.ID) == "" { graph.ID = uuid.NewString() } filter := brandOwnerFilter(graph.TenantID, graph.OwnerUID, graph.BrandID) update := bson.M{ "$set": bson.M{ "seed": graph.Seed, "nodes": graph.Nodes, "edges": graph.Edges, "brave_sources": graph.BraveSources, "expand_strategy": graph.ExpandStrategy, "pain_tag_count": graph.PainTagCount, "generated_at": graph.GeneratedAt, "update_at": graph.UpdateAt, "brand_id": graph.BrandID, }, "$setOnInsert": bson.M{ "_id": graph.ID, "tenant_id": graph.TenantID, "owner_uid": graph.OwnerUID, "create_at": graph.CreateAt, }, } opts := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After) var out entity.Graph err := r.collection.FindOneAndUpdate(ctx, filter, update, opts).Decode(&out) if err != nil { return nil, err } return &out, nil } func (r *mongoRepository) UpsertByTopic(ctx context.Context, graph *entity.Graph) (*entity.Graph, error) { if r.collection == nil { return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured") } if graph == nil || strings.TrimSpace(graph.TopicID) == "" { return nil, app.For(code.Brand).InputMissingRequired("topic_id is required") } now := clock.NowUnixNano() graph.UpdateAt = now if graph.CreateAt == 0 { graph.CreateAt = now } if strings.TrimSpace(graph.ID) == "" { graph.ID = uuid.NewString() } filter := topicOwnerFilter(graph.TenantID, graph.OwnerUID, graph.TopicID) update := bson.M{ "$set": bson.M{ "seed": graph.Seed, "nodes": graph.Nodes, "edges": graph.Edges, "brave_sources": graph.BraveSources, "expand_strategy": graph.ExpandStrategy, "pain_tag_count": graph.PainTagCount, "generated_at": graph.GeneratedAt, "update_at": graph.UpdateAt, "brand_id": graph.BrandID, "topic_id": graph.TopicID, }, "$setOnInsert": bson.M{ "_id": graph.ID, "tenant_id": graph.TenantID, "owner_uid": graph.OwnerUID, "create_at": graph.CreateAt, }, } opts := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After) var out entity.Graph err := r.collection.FindOneAndUpdate(ctx, filter, update, opts).Decode(&out) if err != nil { return nil, err } return &out, nil } func (r *mongoRepository) FindByTopic(ctx context.Context, tenantID, ownerUID, topicID string) (*entity.Graph, error) { if r.collection == nil { return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured") } var out entity.Graph err := r.collection.FindOne(ctx, topicOwnerFilter(tenantID, ownerUID, topicID)).Decode(&out) if err == mongo.ErrNoDocuments { return nil, app.For(code.Brand).ResNotFound("knowledge graph not found") } if err != nil { return nil, err } return &out, nil } func (r *mongoRepository) AttachTopicID(ctx context.Context, tenantID, ownerUID, brandID, topicID string) error { if r.collection == nil { return app.For(code.Brand).DBUnavailable("Mongo is not configured") } _, err := r.collection.UpdateOne( ctx, brandOwnerFilter(tenantID, ownerUID, brandID), bson.M{"$set": bson.M{"topic_id": strings.TrimSpace(topicID)}}, ) return err } func (r *mongoRepository) UpdateNodesByTopic( ctx context.Context, tenantID, ownerUID, topicID string, nodes []libkg.Node, painTagCount int, ) (*entity.Graph, error) { if r.collection == nil { return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured") } now := clock.NowUnixNano() var out entity.Graph err := r.collection.FindOneAndUpdate( ctx, topicOwnerFilter(tenantID, ownerUID, topicID), bson.M{"$set": bson.M{ "nodes": nodes, "pain_tag_count": painTagCount, "update_at": now, }}, options.FindOneAndUpdate().SetReturnDocument(options.After), ).Decode(&out) if err == mongo.ErrNoDocuments { return nil, app.For(code.Brand).ResNotFound("knowledge graph not found") } if err != nil { return nil, err } return &out, nil } func (r *mongoRepository) FindByBrand(ctx context.Context, tenantID, ownerUID, brandID string) (*entity.Graph, error) { if r.collection == nil { return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured") } var out entity.Graph err := r.collection.FindOne(ctx, brandOwnerFilter(tenantID, ownerUID, brandID)).Decode(&out) if err == mongo.ErrNoDocuments { return nil, app.For(code.Brand).ResNotFound("knowledge graph not found") } if err != nil { return nil, err } return &out, nil } func (r *mongoRepository) UpdateNodes( ctx context.Context, tenantID, ownerUID, brandID string, nodes []libkg.Node, painTagCount int, ) (*entity.Graph, error) { if r.collection == nil { return nil, app.For(code.Brand).DBUnavailable("Mongo is not configured") } now := clock.NowUnixNano() var out entity.Graph err := r.collection.FindOneAndUpdate( ctx, brandOwnerFilter(tenantID, ownerUID, brandID), bson.M{"$set": bson.M{ "nodes": nodes, "pain_tag_count": painTagCount, "update_at": now, }}, options.FindOneAndUpdate().SetReturnDocument(options.After), ).Decode(&out) if err == mongo.ErrNoDocuments { return nil, app.For(code.Brand).ResNotFound("knowledge graph not found") } if err != nil { return nil, err } return &out, nil }