import asyncio
import os
import time
import uuid

import reflex as rx


class State(rx.State):
    messages: list[str] = []
    current_pid: int = 0

    @rx.event
    def on_load(self):
        self.current_pid = os.getpid()

    @rx.event(background=True)
    async def many_message_stream(self):
        for _ in range(4):
            yield [
                State.message_stream,
                rx.call_script("socket.disconnect(); window.location.reload();"),
            ]
            await asyncio.sleep(1)

    @rx.event(background=True)
    async def message_stream(self):
        """A background stream that appends messages every 5 seconds."""
        task_instance = uuid.uuid4().hex[:6]
        task_pid = os.getpid()
        task_sid = self.router.session.session_id

        for count in range(1, 10):
            async with self:
                current_sid = self.router.session.session_id
                self.messages.append(
                    f"[{time.time()}] [{task_instance} -> {count}] "
                    f"from {task_pid}:{task_sid} (now {self.current_pid}:{current_sid})"
                )
            await asyncio.sleep(5)

    @rx.event
    def clear_messages(self):
        """Clear all messages."""
        self.messages = []


class CounterState(rx.State):
    count: int = 0
    last_toucher: str = ""

    @rx.event(background=True)
    async def increment(self):
        task_sid = self.router.session.session_id
        task_pid = os.getpid()
        for _ in range(1000):
            async with self:
                self.last_toucher = f"{task_pid}:{task_sid}"
                self.count += 1
           # await asyncio.sleep(0.1)


def index() -> rx.Component:
    return rx.container(
        rx.color_mode.button(position="top-right"),
        rx.vstack(
            rx.hstack(
                rx.button("Start Streamers", on_click=State.many_message_stream, color_scheme="green"),
                rx.button("Single Streamer", on_click=State.message_stream, color_scheme="blue"),
                rx.button(rx.icon("eraser"), on_click=State.clear_messages, color_scheme="red"),
                rx.badge(State.current_pid, color_scheme="purple", font_size="0.8em"),
                rx.badge(State.router.session.client_token, color_scheme="orange", font_size="0.8em"),
                rx.badge(State.router.session.session_id, color_scheme="teal", font_size="0.8em"),
            ),
            rx.hstack(
                rx.button("Increment Counter", on_click=CounterState.increment, color_scheme="green"),
                rx.badge(CounterState.count, color_scheme="purple", font_size="0.8em"),
                rx.badge(CounterState.last_toucher, color_scheme="orange", font_size="0.8em"),
            ),
            rx.foreach(
                State.messages.reverse(),
                lambda msg: rx.text(msg),
            ),
        ),
    )


app = rx.App()
app.add_page(index, on_load=State.on_load, title="Background Stream")