import textwrap
from contextvars import ContextVar
from unittest.mock import MagicMock

import pytest
from pendulum.datetime import DateTime

import prefect.settings
from prefect import flow, task
from prefect.context import (
    GLOBAL_SETTINGS_CONTEXT,
    ContextModel,
    FlowRunContext,
    SettingsContext,
    TagsContext,
    TaskRunContext,
    get_run_context,
    get_settings_context,
    hydrated_context,
    root_settings_context,
    serialize_context,
    tags,
    use_profile,
)
from prefect.exceptions import MissingContextError
from prefect.results import ResultFactory
from prefect.settings import (
    DEFAULT_PROFILES_PATH,
    PREFECT_API_KEY,
    PREFECT_API_URL,
    PREFECT_HOME,
    PREFECT_PROFILES_PATH,
    Profile,
    ProfilesCollection,
    save_profiles,
    temporary_settings,
)
from prefect.task_runners import ThreadPoolTaskRunner


class ExampleContext(ContextModel):
    __var__ = ContextVar("test")

    x: int


def test_context_enforces_types():
    with pytest.raises(ValueError):
        ExampleContext(x="hello")


def test_context_get_outside_context_is_null():
    assert ExampleContext.get() is None


def test_single_context_object_cannot_be_entered_multiple_times():
    context = ExampleContext(x=1)
    with context:
        with pytest.raises(RuntimeError, match="Context already entered"):
            with context:
                pass


def test_copied_context_object_can_be_reentered():
    context = ExampleContext(x=1)
    with context:
        with context.model_copy():
            assert ExampleContext.get().x == 1


def test_exiting_a_context_more_than_entering_raises():
    context = ExampleContext(x=1)

    with pytest.raises(RuntimeError, match="Asymmetric use of context"):
        with context:
            context.__exit__()


def test_context_exit_restores_previous_context():
    with ExampleContext(x=1):
        with ExampleContext(x=2):
            with ExampleContext(x=3):
                assert ExampleContext.get().x == 3
            assert ExampleContext.get().x == 2
        assert ExampleContext.get().x == 1
    assert ExampleContext.get() is None


async def test_flow_run_context(prefect_client):
    @flow
    def foo():
        pass

    test_task_runner = ThreadPoolTaskRunner()
    flow_run = await prefect_client.create_flow_run(foo)
    result_factory = await ResultFactory.from_flow(foo)

    with FlowRunContext(
        flow=foo,
        flow_run=flow_run,
        client=prefect_client,
        task_runner=test_task_runner,
        result_factory=result_factory,
        parameters={"x": "y"},
    ):
        ctx = FlowRunContext.get()
        assert ctx.flow is foo
        assert ctx.flow_run == flow_run
        assert ctx.client is prefect_client
        assert ctx.task_runner is test_task_runner
        assert ctx.result_factory == result_factory
        assert isinstance(ctx.start_time, DateTime)
        assert ctx.parameters == {"x": "y"}


async def test_task_run_context(prefect_client, flow_run):
    @task
    def foo():
        pass

    task_run = await prefect_client.create_task_run(foo, flow_run.id, dynamic_key="")
    result_factory = await ResultFactory.default_factory()

    with TaskRunContext(
        task=foo,
        task_run=task_run,
        client=prefect_client,
        result_factory=result_factory,
        parameters={"foo": "bar"},
    ):
        ctx = TaskRunContext.get()
        assert ctx.task is foo
        assert ctx.task_run == task_run
        assert ctx.result_factory == result_factory
        assert isinstance(ctx.start_time, DateTime)
        assert ctx.parameters == {"foo": "bar"}


@pytest.fixture
def remove_existing_settings_context():
    token = SettingsContext.__var__.set(None)
    try:
        yield
    finally:
        SettingsContext.__var__.reset(token)


async def test_get_run_context(prefect_client, local_filesystem):
    @flow
    def foo():
        pass

    @task
    def bar():
        pass

    test_task_runner = ThreadPoolTaskRunner()
    flow_run = await prefect_client.create_flow_run(foo)
    task_run = await prefect_client.create_task_run(bar, flow_run.id, dynamic_key="")

    with pytest.raises(RuntimeError):
        get_run_context()

    with pytest.raises(MissingContextError):
        get_run_context()

    with FlowRunContext(
        flow=foo,
        flow_run=flow_run,
        client=prefect_client,
        task_runner=test_task_runner,
        result_factory=await ResultFactory.from_flow(foo, client=prefect_client),
        parameters={"x": "y"},
    ) as flow_ctx:
        assert get_run_context() is flow_ctx

        with TaskRunContext(
            task=bar,
            task_run=task_run,
            client=prefect_client,
            result_factory=await ResultFactory.from_task(bar, client=prefect_client),
            parameters={"foo": "bar"},
        ) as task_ctx:
            assert get_run_context() is task_ctx, "Task context takes precedence"

        assert get_run_context() is flow_ctx, "Flow context is restored and retrieved"


class TestSettingsContext:
    @pytest.fixture(autouse=True)
    def temporary_profiles_path(self, tmp_path):
        path = tmp_path / "profiles.toml"
        with temporary_settings(
            updates={PREFECT_HOME: tmp_path, PREFECT_PROFILES_PATH: path}
        ):
            yield path

    def test_settings_context_variable(self):
        with SettingsContext(
            profile=Profile(name="test", settings={}),
            settings=prefect.settings.get_settings_from_env(),
        ) as context:
            assert get_settings_context() is context
            assert context.profile == Profile(name="test", settings={})
            assert context.settings == prefect.settings.get_settings_from_env()

    def test_get_settings_context_missing(self, monkeypatch):
        # It's kind of hard to actually exit the default profile, so we patch `get`
        monkeypatch.setattr(
            "prefect.context.SettingsContext.get", MagicMock(return_value=None)
        )
        with pytest.raises(MissingContextError, match="No settings context found"):
            get_settings_context()

    def test_creates_home(self, tmp_path):
        home = tmp_path / "home"
        assert not home.exists()
        with temporary_settings(updates={PREFECT_HOME: home}):
            pass

        assert home.exists()

    def test_settings_context_uses_settings(self, temporary_profiles_path):
        temporary_profiles_path.write_text(
            textwrap.dedent(
                """
                [profiles.foo]
                PREFECT_API_URL="test"
                """
            )
        )
        with use_profile("foo") as ctx:
            assert prefect.settings.PREFECT_API_URL.value() == "test"
            assert ctx.settings == prefect.settings.get_current_settings()
            assert ctx.profile == Profile(
                name="foo",
                settings={PREFECT_API_URL: "test"},
                source=temporary_profiles_path,
            )

    def test_settings_context_does_not_setup_logging(self, monkeypatch):
        setup_logging = MagicMock()
        monkeypatch.setattr(
            "prefect.logging.configuration.setup_logging", setup_logging
        )
        with use_profile("default"):
            setup_logging.assert_not_called()

    def test_settings_context_nesting(self, temporary_profiles_path):
        temporary_profiles_path.write_text(
            textwrap.dedent(
                """
                [profiles.foo]
                PREFECT_API_URL="foo"

                [profiles.bar]
                PREFECT_API_URL="bar"
                """
            )
        )
        with use_profile("foo") as foo_context:
            with use_profile("bar") as bar_context:
                assert bar_context.settings == prefect.settings.get_current_settings()
                assert (
                    prefect.settings.PREFECT_API_URL.value_from(bar_context.settings)
                    == "bar"
                )
                assert bar_context.profile == Profile(
                    name="bar",
                    settings={PREFECT_API_URL: "bar"},
                    source=temporary_profiles_path,
                )
            assert foo_context.settings == prefect.settings.get_current_settings()
            assert (
                prefect.settings.PREFECT_API_URL.value_from(foo_context.settings)
                == "foo"
            )
            assert foo_context.profile == Profile(
                name="foo",
                settings={PREFECT_API_URL: "foo"},
                source=temporary_profiles_path,
            )

    @pytest.fixture
    def foo_profile(self, temporary_profiles_path):
        profile = Profile(
            name="foo",
            settings={PREFECT_API_KEY: "xxx"},
            source=temporary_profiles_path,
        )
        save_profiles(ProfilesCollection(profiles=[profile]))
        return profile

    def test_root_settings_context_default(self, monkeypatch):
        use_profile = MagicMock()
        monkeypatch.setattr("prefect.context.use_profile", use_profile)
        result = root_settings_context()
        use_profile.assert_called_once_with(
            Profile(name="default", settings={}, source=DEFAULT_PROFILES_PATH),
            override_environment_variables=False,
        )
        use_profile().__enter__.assert_called_once_with()
        assert result is not None

    @pytest.mark.parametrize(
        "cli_command",
        [
            # No profile name provided
            ["prefect", "--profile"],
            # Not called via `prefect` CLI
            ["foobar", "--profile", "test"],
        ],
    )
    def test_root_settings_context_default_if_cli_args_do_not_match_format(
        self, monkeypatch, cli_command
    ):
        use_profile = MagicMock()
        monkeypatch.setattr("prefect.context.use_profile", use_profile)
        monkeypatch.setattr("sys.argv", cli_command)
        result = root_settings_context()
        use_profile.assert_called_once_with(
            Profile(name="default", settings={}, source=DEFAULT_PROFILES_PATH),
            override_environment_variables=False,
        )
        use_profile().__enter__.assert_called_once_with()
        assert result is not None

    def test_root_settings_context_respects_cli(self, monkeypatch, foo_profile):
        use_profile = MagicMock()
        monkeypatch.setattr("prefect.context.use_profile", use_profile)
        monkeypatch.setattr("sys.argv", ["/prefect", "--profile", "foo"])
        result = root_settings_context()
        use_profile.assert_called_once_with(
            foo_profile,
            override_environment_variables=True,
        )
        use_profile().__enter__.assert_called_once_with()
        assert result is not None

    def test_root_settings_context_respects_environment_variable(
        self, monkeypatch, foo_profile
    ):
        use_profile = MagicMock()
        monkeypatch.setattr("prefect.context.use_profile", use_profile)
        monkeypatch.setenv("PREFECT_PROFILE", "foo")
        root_settings_context()
        use_profile.assert_called_once_with(
            foo_profile, override_environment_variables=False
        )

    def test_root_settings_context_missing_cli(self, monkeypatch, capsys):
        use_profile = MagicMock()
        monkeypatch.setattr("prefect.context.use_profile", use_profile)
        monkeypatch.setattr("sys.argv", ["/prefect", "--profile", "bar"])
        root_settings_context()
        _, err = capsys.readouterr()
        assert (
            "profile 'bar' set by command line argument not found. The default profile"
            " will be used instead." in err
        )

    def test_root_settings_context_missing_environment_variables(
        self, monkeypatch, capsys
    ):
        use_profile = MagicMock()
        monkeypatch.setattr("prefect.context.use_profile", use_profile)
        monkeypatch.setenv("PREFECT_PROFILE", "bar")
        root_settings_context()
        _, err = capsys.readouterr()
        assert (
            "profile 'bar' set by environment variable not found. The default profile"
            " will be used instead." in err
        )

    @pytest.mark.usefixtures("remove_existing_settings_context")
    def test_root_settings_context_accessible_in_new_thread(self):
        from concurrent.futures.thread import ThreadPoolExecutor

        with ThreadPoolExecutor() as executor:
            result = executor.submit(get_settings_context).result()

        assert result == GLOBAL_SETTINGS_CONTEXT

    @pytest.mark.usefixtures("remove_existing_settings_context")
    def test_root_settings_context_accessible_in_new_loop(self):
        from anyio.from_thread import start_blocking_portal

        with start_blocking_portal() as portal:
            result = portal.call(get_settings_context)

        assert result == GLOBAL_SETTINGS_CONTEXT


class TestSerializeContext:
    def test_empty(self):
        serialized = serialize_context()
        assert serialized == {
            "flow_run_context": {},
            "task_run_context": {},
            "tags_context": {},
            "settings_context": SettingsContext.get().serialize(),
        }

    async def test_with_flow_run_context(self, prefect_client):
        @flow
        def foo():
            pass

        test_task_runner = ThreadPoolTaskRunner()
        flow_run = await prefect_client.create_flow_run(foo)
        result_factory = await ResultFactory.from_flow(foo)

        with FlowRunContext(
            flow=foo,
            flow_run=flow_run,
            client=prefect_client,
            task_runner=test_task_runner,
            result_factory=result_factory,
            parameters={"x": "y"},
        ) as flow_run_context:
            serialized = serialize_context()
            assert serialized == {
                "flow_run_context": flow_run_context.serialize(),
                "task_run_context": {},
                "tags_context": {},
                "settings_context": SettingsContext.get().serialize(),
            }

    async def test_with_task_run_context(self, prefect_client, flow_run):
        @task
        def bar():
            pass

        task_run = await prefect_client.create_task_run(
            bar, flow_run.id, dynamic_key=""
        )

        with TaskRunContext(
            task=bar,
            task_run=task_run,
            client=prefect_client,
            result_factory=await ResultFactory.from_task(bar, client=prefect_client),
            parameters={"foo": "bar"},
        ) as task_ctx:
            serialized = serialize_context()
            assert serialized == {
                "flow_run_context": {},
                "task_run_context": task_ctx.serialize(),
                "tags_context": {},
                "settings_context": SettingsContext.get().serialize(),
            }

    def test_with_tags_context(self):
        with tags("a", "b") as current_tags:
            serialized = serialize_context()
            assert serialized == {
                "flow_run_context": {},
                "task_run_context": {},
                "tags_context": {"current_tags": current_tags},
                "settings_context": SettingsContext.get().serialize(),
            }

    def test_with_multiple_contexts(self):
        with tags("a", "b") as current_tags:
            with temporary_settings(
                updates={PREFECT_API_KEY: "test", PREFECT_API_URL: "test"}
            ):
                serialized = serialize_context()
                assert serialized == {
                    "flow_run_context": {},
                    "task_run_context": {},
                    "tags_context": {"current_tags": current_tags},
                    "settings_context": SettingsContext.get().serialize(),
                }
                assert (
                    serialized["settings_context"]["settings"]["PREFECT_API_KEY"]
                    == "test"
                )
                assert (
                    serialized["settings_context"]["settings"]["PREFECT_API_URL"]
                    == "test"
                )


class TestHydratedContext:
    def test_empty(self):
        initial_settings_context = SettingsContext.get()
        with hydrated_context({}):
            assert FlowRunContext.get() is None
            assert TaskRunContext.get() is None
            assert TagsContext.get().current_tags == set()
            assert SettingsContext.get() == initial_settings_context

    async def test_with_flow_run_context(self, prefect_client):
        @flow
        def foo():
            pass

        test_task_runner = ThreadPoolTaskRunner()
        flow_run = await prefect_client.create_flow_run(foo)
        result_factory = await ResultFactory.from_flow(foo)
        flow_run_context = FlowRunContext(
            flow=foo,
            flow_run=flow_run,
            client=prefect_client,
            task_runner=test_task_runner,
            result_factory=result_factory,
            parameters={"x": "y"},
        )

        with hydrated_context(
            {
                "flow_run_context": flow_run_context.serialize(),
            }
        ):
            hydrated_flow_run_context = FlowRunContext.get()
            assert hydrated_flow_run_context.flow is foo
            assert hydrated_flow_run_context.flow_run == flow_run
            assert (
                hydrated_flow_run_context.client is not None
            )  # this won't be the same object as the original client
            assert (
                hydrated_flow_run_context.task_runner is not None
            )  # this won't be the same object as the original task runner
            assert (
                hydrated_flow_run_context.result_factory is not None
            )  # this won't be the same object as the original result factory
            assert isinstance(hydrated_flow_run_context.start_time, DateTime)
            assert hydrated_flow_run_context.parameters == {"x": "y"}

    async def test_with_task_run_context(self, prefect_client, flow_run):
        @task
        def bar():
            pass

        task_run = await prefect_client.create_task_run(
            bar, flow_run.id, dynamic_key=""
        )
        task_ctx = TaskRunContext(
            task=bar,
            task_run=task_run,
            client=prefect_client,
            result_factory=await ResultFactory.from_task(bar, client=prefect_client),
            parameters={"foo": "bar"},
        )

        with hydrated_context(
            {
                "task_run_context": task_ctx.serialize(),
            }
        ):
            hydrated_task_ctx = TaskRunContext.get()
            assert hydrated_task_ctx.task is bar
            assert hydrated_task_ctx.task_run == task_run
            assert hydrated_task_ctx.result_factory is not None
            assert isinstance(hydrated_task_ctx.start_time, DateTime)
            assert hydrated_task_ctx.parameters == {"foo": "bar"}

    def test_with_tags_context(self):
        with hydrated_context(
            {
                "tags_context": {"current_tags": {"a", "b"}},
            }
        ):
            assert TagsContext.get().current_tags == {"a", "b"}

    def test_with_settings_context(self):
        with hydrated_context(
            {
                "settings_context": {
                    "profile": {"name": "test", "settings": {}, "source": None},
                    "settings": {"PREFECT_API_KEY": "test", "PREFECT_API_URL": "test"},
                },
            }
        ):
            assert SettingsContext.get().profile == Profile(
                name="test",
                settings={},
                source=None,
            )
            assert SettingsContext.get().settings.PREFECT_API_KEY == "test"
            assert SettingsContext.get().settings.PREFECT_API_URL == "test"
