# Copyright (C) 2005-2025 Splunk Inc. All Rights Reserved.

import time

from ITOA.setup_logging import getLogger
from feature_flagging.state.state import FeatureFlaggingState
from ITOA.storage.itoa_storage import ITOAStorage


class FeatureFlaggingStateStore(object):
    KEY = 'suite_state'
    COLLECTION = 'itsi_feature_flagging_state'
    OBJECT_TYPE = 'suite_state'

    def __init__(self, session_key):
        self.logger = getLogger(logger_name='itsi.feature_flagging.FeatureFlaggingStateManager')
        self.session_key = session_key
        self.kvstore = ITOAStorage(collection=self.COLLECTION)
        assert self.kvstore.wait_for_storage_init(self.session_key), "Storage is not initialized after wait."

    def set_state(self, state_obj):
        """
        This method stores the FeatureFlaggingState state_obj into the KV Store
        """
        state_data = state_obj.to_json()
        state_data['_key'] = self.KEY
        state_data['object_type'] = self.OBJECT_TYPE
        state_data['mod_time'] = time.time()
        # batch_save essentially does upsert
        self.kvstore.batch_save(self.session_key, 'nobody', [state_data])

    def get_state(self):
        """
        This method retrieves the FeatureFlaggingState state_obj from the KV Store
        """
        state_dict = self.kvstore.get(self.session_key, 'nobody', self.COLLECTION, self.KEY)

        if state_dict is not None:
            # remove KV store related fields
            state_dict.pop('_user', None)
            state_dict.pop('_key', None)
            state_dict.pop('object_type', None)
            state_dict.pop('mod_time', None)

            return FeatureFlaggingState.from_dict(state_dict)

        return None