"""Repro for dirty substate missing from parent state.

Must use StateManagerRedis!

1. Cause common parent state to be touched.
    1a. With oplock enabled, this can be a normal parent var update in a separate `async with self`
    1b. Without oplock, this can be any parent backend var update in the same `async with self`
2. Update a sibling state's backend var to mark it dirty, but NOT result in a delta being sent.
    This causes the parent's dirty_substates to be persisted to redis, because
    `App.modify_state` only calls `_clean()` when a delta is produced.
3. If oplock is enabled, wait for the lock expiration.
4. Enter a new `async with self` and make any change.
    Because the parent state is loaded from redis with dirty_substates already
    set, but the actual substates are not loaded resulting in an error during
    delta resolution.
"""

import asyncio
import reflex as rx


class State(rx.State):
    """The app state."""
    count: int = 0
    _backend_base: int = 0


class BrotherState(State):
    """A sibling state to test oplock missing substate issue."""
    brother_count: int = 0

    @rx.event(background=True)
    async def mark_dirty_no_delta_oplock(self):
        from reflex.istate.manager.redis import StateManagerRedis

        async with self:
            # Must touch any var in parent state in a separate context, so it is marked as touched.
            self.count += 1
        async with self:
            # Must update another state's backend var to mark it dirty, but get no delta.
            sister_state = await self.get_state(SisterState)
            sister_state._backend_sis += 1
            parent_state = self.parent_state
            print("parent dirty_substates 1:", self.parent_state.dirty_substates)
        print(f"parent _was_touched: {parent_state._was_touched} {sister_state._was_touched}")
        # Must wait for the oplock lease to expire.
        state_manager = rx.state.get_state_manager()
        if isinstance(state_manager, StateManagerRedis):
            await asyncio.sleep(state_manager.lock_expiration / 1000)
        async with self:
            # Now when refetching the common parent from redis, it should still have dirty_substates, but the actual substates will not be fetched.
            print("parent dirty_substates 2:", self.parent_state.dirty_substates)
            self.brother_count += 1

    @rx.event(background=True)
    async def mark_dirty_no_delta(self):
        async with self:
            # Must touch backend var in parent state, so it will get persisted.
            self._backend_base += 1
            # Must update another state's backend var to mark it dirty, but get no delta.
            sister_state = await self.get_state(SisterState)
            sister_state._backend_sis += 1
            parent_state = self.parent_state
            print("parent dirty_substates 1:", self.parent_state.dirty_substates)
        print(f"parent _was_touched: {parent_state._was_touched} {sister_state._was_touched}")
        async with self:
            # Now when refetching the common parent from redis, it should still have dirty_substates, but the actual substates will not be fetched.
            print("parent dirty_substates 2:", self.parent_state.dirty_substates)
            self.brother_count += 1
        async with self:
            # No changes here, should not call _clean
            print(self.brother_count)


class SisterState(State):
    """Another sibling state to test oplock missing substate issue."""
    _backend_sis: int = 0


def index() -> rx.Component:
    return rx.container(
        rx.hstack(
            rx.button(
                "Repro w/ Oplock",
                on_click=BrotherState.mark_dirty_no_delta_oplock,
            ),
            rx.button(
                "Repro w/o Oplock",
                on_click=BrotherState.mark_dirty_no_delta,
            ),
        ),
    )


app = rx.App()
app.add_page(index)