package storage

import (
	"context"
	"fmt"

	"cloud.google.com/go/firestore"
	"git.pengzhan.dev/aimaren/internal/crawler"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	StateCollection   = "hermes_state"
	AppStateDocument  = "main"
	UserStateDocument = "users"
)

// FirestoreClient implements Storer
type FirestoreClient struct {
	client *firestore.Client
}

func NewFirestoreClient(ctx context.Context, projectID string) (*FirestoreClient, error) {
	client, err := firestore.NewClient(ctx, projectID)
	if err != nil {
		return nil, err
	}
	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(AppStateDocument).Get(ctx)
	if err != nil {
		if status.Code(err) == codes.NotFound {
			return &AppState{
				Bags:    make(map[string]crawler.Bag),
				ChatIDs: []int64{},
			}, nil
		}
		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(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(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
}

// 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
}