from contextvars import ContextVar
from typing import Any

import pendulum
from pydantic import Field

from prefect import flow, task
from prefect.context import ContextModel


class SharedState(ContextModel):
    """
    A shared state model for maintaining context across a flow and its tasks.
    """

    __var__: ContextVar = ContextVar("shared_state")

    account_id: str
    other_data: dict[str, Any] = Field(default_factory=dict)
    start_time: pendulum.DateTime = Field(default_factory=lambda: pendulum.now("UTC"))


@task
def task1():
    state = SharedState.get()
    if state:
        print(f"Task 1 executing for account {state.account_id}")
        print(f"Other data: {state.other_data}")
    else:
        print("No shared state found in Task 1")


@task
def task2():
    state = SharedState.get()
    if state:
        print(f"Task 2 executing for account {state.account_id}")
        print(f"Time since start: {pendulum.now('UTC') - state.start_time}")
    else:
        print("No shared state found in Task 2")


@flow
def main_flow(shared_state: SharedState):
    with shared_state:
        task1()
        task2()


if __name__ == "__main__":
    main_flow(shared_state=dict(account_id="12345", other_data={"key": "value"}))