feat: Scaffolding bot

internal/storage/firestore.go

@@ -2,6 +2,7 @@ package storage
 
 import (
 	"context"
+	"fmt"
 
 	"cloud.google.com/go/firestore"
 	"git.pengzhan.dev/aimaren/internal/crawler"
@@ -10,10 +11,12 @@ import (
 )
 
 const (
-	stateCollection = "hermes_state"
-	stateDocument   = "main"
+	StateCollection   = "hermes_state"
+	AppStateDocument  = "main"
+	UserStateDocument = "users"
 )
 
+// FirestoreClient implements Storer
 type FirestoreClient struct {
 	client *firestore.Client
 }
@@ -26,11 +29,14 @@ func NewFirestoreClient(ctx context.Context, projectID string) (*FirestoreClient
 	return &FirestoreClient{client: client}, nil
 }
 
+func (fs *FirestoreClient) Close() {
+	fs.client.Close()
+}
+
 // FetchAppState retrieves the entire application state from a single document.
 func (fs *FirestoreClient) FetchAppState(ctx context.Context) (*AppState, error) {
-	doc, err := fs.client.Collection(stateCollection).Doc(stateDocument).Get(ctx)
+	doc, err := fs.client.Collection(StateCollection).Doc(AppStateDocument).Get(ctx)
 	if err != nil {
-		// If the doc doesn't exist, return a new, empty AppState.
 		if status.Code(err) == codes.NotFound {
 			return &AppState{
 				Bags:    make(map[string]crawler.Bag),
@@ -39,29 +45,59 @@ func (fs *FirestoreClient) FetchAppState(ctx context.Context) (*AppState, error)
 		}
 		return nil, err
 	}
-
 	var state AppState
 	if err := doc.DataTo(&state); err != nil {
 		return nil, err
 	}
+	if state.Bags == nil {
+		state.Bags = make(map[string]crawler.Bag)
+	}
+	if state.ChatIDs == nil {
+		state.ChatIDs = []int64{}
+	}
 	return &state, nil
 }
 
 // UpdateAppState writes the entire application state back to the document.
 func (fs *FirestoreClient) UpdateAppState(ctx context.Context, newState *AppState) error {
-	_, err := fs.client.Collection(stateCollection).Doc(stateDocument).Set(ctx, newState)
+	_, err := fs.client.Collection(StateCollection).Doc(AppStateDocument).Set(ctx, newState)
 	return err
 }
 
 // AddChatID atomically adds a new chat ID to the list in the main state document.
 func (fs *FirestoreClient) AddChatID(ctx context.Context, chatID int64) error {
-	docRef := fs.client.Collection(stateCollection).Doc(stateDocument)
+	docRef := fs.client.Collection(StateCollection).Doc(AppStateDocument)
 	_, err := docRef.Update(ctx, []firestore.Update{
 		{Path: "chat_ids", Value: firestore.ArrayUnion(chatID)},
 	})
+	if status.Code(err) == codes.NotFound {
+		return fs.UpdateAppState(ctx, &AppState{
+			Bags:    make(map[string]crawler.Bag),
+			ChatIDs: []int64{chatID},
+		})
+	}
 	return err
 }
 
-func (fs *FirestoreClient) Close() {
-	fs.client.Close()
+// FetchUserState retrieves the user state document and converts string keys back to int64.
+func (fs *FirestoreClient) FetchUserState(ctx context.Context) (*UserState, error) {
+	docSnap, err := fs.client.Collection(StateCollection).Doc(UserStateDocument).Get(ctx)
+	if err != nil {
+		if status.Code(err) == codes.NotFound {
+			return &UserState{Users: make(map[string]ChatState)}, nil
+		}
+		return nil, err
+	}
+	var result UserState
+	if err := docSnap.DataTo(&result); err != nil {
+		return nil, fmt.Errorf("failed to decode user state: %w", err)
+	}
+	return &result, nil
+}
+
+// UpdateUserState writes the whole UserState back to Firestore.
+// It converts int64 keys to string keys for Firestore compatibility.
+func (fs *FirestoreClient) UpdateUserState(ctx context.Context, newUserState *UserState) error {
+	_, err := fs.client.Collection(StateCollection).Doc(UserStateDocument).Set(ctx, newUserState)
+	return err
 }