import hashlib
import random
import tempfile

from burr.core import ApplicationBuilder, State, action
from burr.core.action import create_action
from burr.core.parallelism import MapStates
from burr.tracking.client import LocalTrackingClient


@action(reads=[], writes=["x"])
def pick(state: State) -> State:
    return state.update(x=random.random())


@action(reads=[], writes=[])
def back(state: State) -> State:
    return state


class Fan(MapStates):
    def action(self, state, inputs):
        return create_action(pick, "pick")

    def states(self, state, context, inputs):
        for _ in range(3):
            yield state

    def reduce(self, state, results):
        return state.update(xs=[s["x"] for s in results])

    @property
    def reads(self):
        return []

    @property
    def writes(self):
        return ["xs"]


class FixedFan(Fan):
    """Workaround: mix context.sequence_id into the sub-app id so the per-
    invocation sub-apps get distinct ids and the cascaded initializer can't
    collide them onto a prior call's persisted state."""

    def tasks(self, state, context, inputs):
        for task in super().tasks(state, context, inputs):
            task.application_id = hashlib.sha256(
                f"{task.application_id}:{context.sequence_id}".encode()
            ).hexdigest()
            yield task


def run(fan: MapStates) -> list[list[float]]:
    tracker = LocalTrackingClient(project="test", storage_dir=tempfile.mkdtemp())
    app = (
        ApplicationBuilder()
        .with_actions(fan=fan, back=back)
        .with_transitions(("fan", "back"), ("back", "fan"))
        .with_tracker(tracker)
        .initialize_from(
            tracker,
            resume_at_next_action=True,
            default_state={},
            default_entrypoint="fan",
        )
        .build()
    )
    results = []
    for _ in range(3):
        app.run(halt_after=["fan"])
        results.append(list(app.state["xs"]))
    return results


def main() -> int:
    rc = 0
    for label, fan in [("buggy Fan", Fan()), ("FixedFan", FixedFan())]:
        results = run(fan)
        print(f"\n[{label}]")
        for i, xs in enumerate(results):
            print(f"  invocation {i}: {xs}")
        replayed = all(xs == results[0] for xs in results[1:])
        print(f"  -> {'STALE REPLAY' if replayed else 'fresh outputs'}")
        if (label == "buggy Fan") != replayed:
            rc = 1
    return rc


if __name__ == "__main__":
    raise SystemExit(main())