import reflex as rx


class MySharedThing(rx.SharedState):
    my_counter: int = 0

    @rx.event
    async def toggle_link(self):
        if not self._linked_to:
            await self._link_to(await self.get_var_value(State.shared_token))
        else:
            return await self._unlink()

    @rx.event
    def increment(self):
        self.my_counter += 1

    @rx.event
    def decrement(self):
        self.my_counter -= 1

    @rx.var
    def linked_to(self) -> str:
        return self._linked_to or "not linked"

    @rx.var
    def linked_from(self) -> str:
        return ", ".join(self._linked_from) or "no links"

    @rx.event(background=True)
    async def delayed_multi_increment(self, amount: int):
        import asyncio

        for _ in range(amount):
            await asyncio.sleep(1)
            async with self:
                self.my_counter += 1


class State(rx.State):
    @rx.var
    def shared_token(self) -> str:
        return (self.room or "shared_global").replace("_", "-")

    @rx.var
    async def current_count(self) -> int:
        shared_state = await self.get_state(MySharedThing)
        return shared_state.my_counter

    @rx.event
    async def print_current_count(self):
        shared_state = await self.get_state(MySharedThing)
        print(f"Current count is: {shared_state.my_counter}")


def index() -> rx.Component:
    return rx.container(
        rx.color_mode.button(position="top-right"),
        rx.vstack(
            rx.text(f"Shared token: {State.shared_token}"),
            rx.button(f"Linked To: {MySharedThing.linked_to}", on_click=MySharedThing.toggle_link),
            rx.text(f"Linked From: {MySharedThing.linked_from}"),
            rx.heading(State.current_count),
            rx.button(
                "Increment",
                on_click=MySharedThing.increment,
            ),
            rx.button(
                "Increment 5 times with 1s delay",
                on_click=MySharedThing.delayed_multi_increment(5),
            ),
            rx.button(
                "Decrement",
                on_click=MySharedThing.decrement,
            ),
            rx.button(
                "Print Current Count to Console",
                on_click=State.print_current_count,
            ),
        ),
    )


app = rx.App()
app.add_page(index, route="/[room]")
app.add_page(index)