import inspect
import json
import re
import warnings
from datetime import datetime

import dagster as dg
import pytest
from dagster._core.errors import ScheduleExecutionError
from dagster._time import get_current_datetime
from dagster._utils.merger import merge_dicts

# This file tests a lot of parameter name stuff, so these warnings are spurious


def test_scheduler():
    def define_schedules():
        return [
            dg.ScheduleDefinition(
                name="my_schedule",
                cron_schedule="* * * * *",
                job_name="test_job",
                run_config={},
            )
        ]

    @dg.schedule(cron_schedule="* * * * *", job_name="foo_job")
    def echo_time_schedule(context):
        return {
            "echo_time": (
                context.scheduled_execution_time.isoformat()
                if context.scheduled_execution_time
                else ""
            )
        }

    @dg.schedule(
        cron_schedule="* * * * *",
        job_name="foo_job",
        should_execute=lambda x: False,
    )
    def always_skip_schedule():
        return {}

    context_without_time = dg.build_schedule_context()

    execution_time = datetime(year=2019, month=2, day=27)

    context_with_time = dg.build_schedule_context(scheduled_execution_time=execution_time)

    with pytest.raises(ScheduleExecutionError):
        echo_time_schedule.evaluate_tick(context_without_time)

    execution_data = echo_time_schedule.evaluate_tick(context_with_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1
    assert execution_data.run_requests[0].run_config == {"echo_time": execution_time.isoformat()}

    execution_data = always_skip_schedule.evaluate_tick(context_with_time)
    assert execution_data.skip_message
    assert (
        execution_data.skip_message
        == "should_execute function for always_skip_schedule returned false."
    )


def test_schedule_decorators_sanity():
    @dg.op
    def do_nothing(_):
        pass

    @dg.job
    def foo_job():
        do_nothing()

    @dg.schedule(cron_schedule="* * * * *", job_name="foo_job")
    def foo_schedule():
        """Fake doc block."""
        return {}

    # Ensure that schedule definition inherits properties from wrapped fxn
    assert foo_schedule.__doc__ == """Fake doc block."""

    assert not foo_schedule.execution_timezone

    @dg.schedule(
        cron_schedule="* * * * *",
        job_name="foo_job",
        execution_timezone="US/Central",
    )
    def foo_schedule_timezone():
        return {}

    assert foo_schedule_timezone.execution_timezone == "US/Central"

    with pytest.raises(
        dg.DagsterInvalidDefinitionError,
        match=re.escape(
            "Invalid execution timezone MadeUpTimeZone for invalid_timezone_foo_schedule"
        ),
    ):

        @dg.schedule(
            cron_schedule="* * * * *",
            job_name="foo_job",
            execution_timezone="MadeUpTimeZone",
        )
        def invalid_timezone_foo_schedule():
            return {}


def _check_partitions(
    partition_schedule_def,
    expected_num_partitions,
    expected_start_date,
    expected_format,
    expected_relative_delta,
):
    partitions = partition_schedule_def.get_partition_set().get_partitions()

    assert len(partitions) == expected_num_partitions

    assert partitions[0].value == expected_start_date
    assert partitions[0].name == expected_start_date.strftime(expected_format)

    for index, partition in enumerate(partitions):
        partition_value = partitions[0].value + (index * expected_relative_delta)
        assert partition.value == partitions[0].value + (index * expected_relative_delta)
        assert partition.name == partition_value.strftime(expected_format)


HOURS_UNTIL_FEBRUARY_27 = 24 * (31 + 26)


def test_schedule_decorators_bad():
    @dg.op
    def do_nothing(_):
        pass

    @dg.job
    def foo_job():
        do_nothing()

    with pytest.raises(dg.DagsterInvalidDefinitionError, match="invalid cron schedule"):

        @dg.schedule(cron_schedule="", job_name="foo_job")
        def bad_cron_string(context):
            return {}

    with pytest.raises(dg.DagsterInvalidDefinitionError, match="invalid cron schedule"):

        @dg.schedule(cron_schedule="bad_schedule_two", job_name="foo_job")
        def bad_cron_string_two(context):
            return {}

    with pytest.raises(dg.DagsterInvalidDefinitionError, match="invalid cron schedule"):

        @dg.schedule(cron_schedule="* * * * * *", job_name="foo_job")
        def bad_cron_string_three(context):
            return {}


def test_schedule_with_nested_tags():
    nested_tags = {"foo": {"bar": "baz"}}

    @dg.schedule(cron_schedule="* * * * *", job_name="foo_job", tags=nested_tags)  # pyright: ignore[reportArgumentType]
    def my_tag_schedule():
        return {}

    assert my_tag_schedule.evaluate_tick(  # pyright: ignore[reportOptionalSubscript]
        dg.build_schedule_context(scheduled_execution_time=get_current_datetime())
    )[0][0].tags == merge_dicts(
        {key: json.dumps(val) for key, val in nested_tags.items()},
        {"dagster/schedule_name": "my_tag_schedule"},
    )


def test_invalid_tag_keys():
    tags = {"my_tag&": "yes", "my_tag#": "yes"}

    # turn off any outer warnings filters, e.g. ignores that are set in pyproject.toml
    warnings.resetwarnings()
    with warnings.catch_warnings(record=True) as caught_warnings:

        @dg.schedule(cron_schedule="* * * * *", job_name="foo_job", tags=tags)
        def my_tag_schedule():
            return {}

        assert len(caught_warnings) == 2
        warning = caught_warnings[0]
        assert "Non-compliant tag keys like ['my_tag&', 'my_tag#'] are deprecated" in str(
            warning.message
        )
        assert warning.filename.endswith("test_schedule.py")

    assert my_tag_schedule.evaluate_tick(  # pyright: ignore[reportOptionalSubscript]
        dg.build_schedule_context(scheduled_execution_time=get_current_datetime())
    )[0][0].tags == merge_dicts(tags, {"dagster/schedule_name": "my_tag_schedule"})


def test_scheduled_jobs():
    @dg.op(config_schema={"foo": dg.Field(dg.String)})
    def foo_op(context):
        pass

    DEFAULT_FOO_CONFIG = {"ops": {"foo_op": {"config": {"foo": "bar"}}}}

    @dg.job(config=DEFAULT_FOO_CONFIG)
    def foo_job():
        foo_op()

    my_schedule = dg.ScheduleDefinition(name="my_schedule", cron_schedule="* * * * *", job=foo_job)

    context_without_time = dg.build_schedule_context()
    execution_data = my_schedule.evaluate_tick(context_without_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1

    dg.validate_run_config(foo_job, execution_data.run_requests[0].run_config)


def test_request_based_schedule():
    context_without_time = dg.build_schedule_context()

    @dg.op(config_schema={"foo": dg.Field(dg.String)})
    def foo_op(context):
        pass

    @dg.job
    def foo_job():
        foo_op()

    FOO_CONFIG = {"ops": {"foo_op": {"config": {"foo": "bar"}}}}

    @dg.schedule(
        cron_schedule="* * * * *",
        job=foo_job,
    )
    def foo_schedule(context):
        return dg.RunRequest(run_key=None, run_config=FOO_CONFIG, tags={"foo": "FOO"})

    # evaluate tick
    execution_data = foo_schedule.evaluate_tick(context_without_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1
    run_request = execution_data.run_requests[0]
    assert run_request.run_config == FOO_CONFIG
    assert run_request.tags.get("foo") == "FOO"

    # test direct invocation
    run_request = foo_schedule(context_without_time)
    assert run_request.run_config == FOO_CONFIG  # pyright: ignore[reportOptionalMemberAccess,reportAttributeAccessIssue]
    assert run_request.tags.get("foo") == "FOO"  # pyright: ignore[reportOptionalMemberAccess,reportAttributeAccessIssue]


def test_request_based_schedule_no_context():
    context_without_time = dg.build_schedule_context()

    @dg.op(config_schema={"foo": dg.Field(dg.String)})
    def foo_op(context):
        pass

    @dg.job
    def foo_job():
        foo_op()

    FOO_CONFIG = {"ops": {"foo_op": {"config": {"foo": "bar"}}}}

    @dg.schedule(
        cron_schedule="* * * * *",
        job=foo_job,
    )
    def foo_schedule():
        return dg.RunRequest(run_key=None, run_config=FOO_CONFIG, tags={"foo": "FOO"})

    # evaluate tick
    execution_data = foo_schedule.evaluate_tick(context_without_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1
    run_request = execution_data.run_requests[0]
    assert run_request.run_config == FOO_CONFIG
    assert run_request.tags.get("foo") == "FOO"

    # test direct invocation
    run_request = foo_schedule()
    assert run_request.run_config == FOO_CONFIG  # pyright: ignore[reportOptionalMemberAccess,reportAttributeAccessIssue]
    assert run_request.tags.get("foo") == "FOO"  # pyright: ignore[reportOptionalMemberAccess,reportAttributeAccessIssue]


def test_config_based_schedule():
    context_without_time = dg.build_schedule_context()

    @dg.op(config_schema={"foo": dg.Field(dg.String)})
    def foo_op(context):
        pass

    @dg.job
    def foo_job():
        foo_op()

    FOO_CONFIG = {"ops": {"foo_op": {"config": {"foo": "bar"}}}}

    @dg.schedule(
        cron_schedule="* * * * *",
        job=foo_job,
        tags={"foo": "FOO"},
    )
    def foo_schedule(context):
        return FOO_CONFIG

    # evaluate tick
    execution_data = foo_schedule.evaluate_tick(context_without_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1
    run_request = execution_data.run_requests[0]
    assert run_request.run_config == FOO_CONFIG
    assert run_request.tags.get("foo") == "FOO"

    # direct invocation
    run_config = foo_schedule(context_without_time)
    assert run_config == FOO_CONFIG


def test_config_based_schedule_no_context():
    context_without_time = dg.build_schedule_context()

    @dg.op(config_schema={"foo": dg.Field(dg.String)})
    def foo_op(context):
        pass

    @dg.job
    def foo_job():
        foo_op()

    FOO_CONFIG = {"ops": {"foo_op": {"config": {"foo": "bar"}}}}

    @dg.schedule(
        cron_schedule="* * * * *",
        job=foo_job,
        tags={"foo": "FOO"},
    )
    def foo_schedule():
        return FOO_CONFIG

    # evaluate tick
    execution_data = foo_schedule.evaluate_tick(context_without_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1
    run_request = execution_data.run_requests[0]
    assert run_request.run_config == FOO_CONFIG
    assert run_request.tags.get("foo") == "FOO"

    # direct invocation
    run_config = foo_schedule()
    assert run_config == FOO_CONFIG


def test_request_based_schedule_generator():
    context_without_time = dg.build_schedule_context()

    @dg.op(config_schema={"foo": dg.Field(dg.String)})
    def foo_op(context):
        pass

    @dg.job
    def foo_job():
        foo_op()

    FOO_CONFIG = {"ops": {"foo_op": {"config": {"foo": "bar"}}}}

    @dg.schedule(
        cron_schedule="* * * * *",
        job=foo_job,
    )
    def foo_schedule(_context):
        yield dg.RunRequest(run_key=None, run_config=FOO_CONFIG, tags={"foo": "FOO"})

    # evaluate tick
    execution_data = foo_schedule.evaluate_tick(context_without_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1
    assert execution_data.run_requests[0].run_config == FOO_CONFIG
    assert execution_data.run_requests[0].tags.get("foo") == "FOO"

    # test direct invocation
    request_generator = foo_schedule(context_without_time)
    assert inspect.isgenerator(request_generator)
    requests = list(request_generator)
    assert len(requests) == 1
    assert requests[0].run_config == FOO_CONFIG
    assert requests[0].tags.get("foo") == "FOO"


def test_request_based_schedule_generator_no_context():
    context_without_time = dg.build_schedule_context()

    @dg.op(config_schema={"foo": dg.Field(dg.String)})
    def foo_op(context):
        pass

    @dg.job
    def foo_job():
        foo_op()

    FOO_CONFIG = {"ops": {"foo_op": {"config": {"foo": "bar"}}}}

    @dg.schedule(
        cron_schedule="* * * * *",
        job=foo_job,
    )
    def foo_schedule():
        yield dg.RunRequest(run_key=None, run_config=FOO_CONFIG, tags={"foo": "FOO"})

    # evaluate tick
    execution_data = foo_schedule.evaluate_tick(context_without_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1
    assert execution_data.run_requests[0].run_config == FOO_CONFIG
    assert execution_data.run_requests[0].tags.get("foo") == "FOO"

    # test direct invocation
    request_generator = foo_schedule()
    assert inspect.isgenerator(request_generator)
    requests = list(request_generator)
    assert len(requests) == 1
    assert requests[0].run_config == FOO_CONFIG
    assert requests[0].tags.get("foo") == "FOO"


def test_vixie_cronstring_schedule():
    context_without_time = dg.build_schedule_context()

    @dg.op
    def foo_op(context):
        pass

    @dg.job
    def foo_job():
        foo_op()

    @dg.schedule(cron_schedule="@daily", job=foo_job)
    def foo_schedule():
        yield dg.RunRequest(run_key=None, run_config={}, tags={"foo": "FOO"})

    # evaluate tick
    execution_data = foo_schedule.evaluate_tick(context_without_time)
    assert execution_data.run_requests
    assert len(execution_data.run_requests) == 1
    assert execution_data.run_requests[0].tags.get("foo") == "FOO"
