from __future__ import print_function

import os
import string
from contextlib import contextmanager

import mock
import pytest
from click.testing import CliRunner

from dagster import (
    PartitionSetDefinition,
    PresetDefinition,
    ScheduleDefinition,
    Shape,
    check,
    lambda_solid,
    pipeline,
    repository,
    seven,
    solid,
)
from dagster.cli.run import run_list_command, run_wipe_command
from dagster.core.host_representation import ExternalPipeline
from dagster.core.instance import DagsterInstance
from dagster.core.launcher import RunLauncher
from dagster.core.storage.pipeline_run import PipelineRun
from dagster.core.test_utils import instance_for_test_tempdir
from dagster.core.types.loadable_target_origin import LoadableTargetOrigin
from dagster.grpc.server import GrpcServerProcess
from dagster.serdes import ConfigurableClass
from dagster.utils import file_relative_path, merge_dicts


def no_print(_):
    return None


@lambda_solid
def do_something():
    return 1


@lambda_solid
def do_input(x):
    return x


@pipeline(
    name='foo', preset_defs=[PresetDefinition(name="test", tags={'foo': 'bar'}),],
)
def foo_pipeline():
    do_input(do_something())


def define_foo_pipeline():
    return foo_pipeline


@pipeline(name='baz', description='Not much tbh')
def baz_pipeline():
    do_input()


def not_a_repo_or_pipeline_fn():
    return 'kdjfkjdf'


not_a_repo_or_pipeline = 123


def define_bar_schedules():
    return {
        'foo_schedule': ScheduleDefinition(
            "foo_schedule", cron_schedule="* * * * *", pipeline_name="test_pipeline", run_config={},
        ),
    }


def define_baz_partitions():
    return {
        'baz_partitions': PartitionSetDefinition(
            name='baz_partitions',
            pipeline_name='baz',
            partition_fn=lambda: string.digits,
            run_config_fn_for_partition=lambda partition: {
                'solids': {'do_input': {'inputs': {'x': {'value': partition.value}}}}
            },
        )
    }


@repository
def bar():
    return {
        'pipelines': {'foo': foo_pipeline, 'baz': baz_pipeline},
        'schedules': define_bar_schedules(),
        'partition_sets': define_baz_partitions(),
    }


@solid
def spew(context):
    context.log.info('HELLO WORLD')


@solid
def fail(context):
    raise Exception('I AM SUPPOSED TO FAIL')


@pipeline
def stdout_pipeline():
    spew()


@pipeline
def stderr_pipeline():
    fail()


class InMemoryRunLauncher(RunLauncher, ConfigurableClass):
    def __init__(self, inst_data=None):
        self._inst_data = inst_data
        self._queue = []

    def launch_run(self, instance, run, external_pipeline):
        check.inst_param(instance, 'instance', DagsterInstance)
        check.inst_param(run, 'run', PipelineRun)
        check.inst_param(external_pipeline, 'external_pipeline', ExternalPipeline)
        self._queue.append(run)
        return run

    def queue(self):
        return self._queue

    @classmethod
    def config_type(cls):
        return Shape({})

    @classmethod
    def from_config_value(cls, inst_data, config_value):
        return cls(inst_data=inst_data,)

    @property
    def inst_data(self):
        return self._inst_data

    def can_terminate(self, run_id):
        return False

    def terminate(self, run_id):
        check.not_implemented('Termintation not supported')


@contextmanager
def default_instance_tempdir(temp_dir, overrides=None):
    default_overrides = {
        'run_launcher': {
            'module': 'dagster_tests.cli_tests.command_tests.test_cli_commands',
            'class': 'InMemoryRunLauncher',
        }
    }
    with instance_for_test_tempdir(
        temp_dir, overrides=merge_dicts(default_overrides, (overrides if overrides else {}))
    ) as instance:
        with mock.patch('dagster.core.instance.DagsterInstance.get') as _instance:
            _instance.return_value = instance
            yield instance


@contextmanager
def default_instance(overrides=None):
    with seven.TemporaryDirectory() as temp_dir:
        with default_instance_tempdir(temp_dir, overrides) as instance:
            yield instance


def managed_grpc_instance():
    return default_instance(overrides={"opt_in": {"local_servers": True}})


@contextmanager
def args_with_instance(gen_instance, *args):
    with gen_instance as instance:
        yield args + (instance,)


def args_with_default_instance(*args):
    return args_with_instance(default_instance(), *args)


def args_with_managed_grpc_instance(*args):
    return args_with_instance(managed_grpc_instance(), *args)


@contextmanager
def grpc_server_bar_kwargs(pipeline_name=None):
    server_process = GrpcServerProcess(
        loadable_target_origin=LoadableTargetOrigin(
            python_file=file_relative_path(__file__, 'test_cli_commands.py'), attribute='bar'
        ),
    )
    with server_process.create_ephemeral_client() as client:
        args = {'grpc_host': client.host}
        if pipeline_name:
            args['pipeline'] = 'foo'
        if client.port:
            args['grpc_port'] = client.port
        if client.socket:
            args['grpc_socket'] = client.socket
        yield args
    server_process.wait()


@contextmanager
def python_bar_cli_args(pipeline_name=None):
    args = [
        '-m',
        'dagster_tests.cli_tests.command_tests.test_cli_commands',
        '-a',
        'bar',
    ]
    if pipeline_name:
        args.append('-p')
        args.append(pipeline_name)
    yield args


@contextmanager
def grpc_server_bar_cli_args(pipeline_name=None):
    server_process = GrpcServerProcess(
        loadable_target_origin=LoadableTargetOrigin(
            python_file=file_relative_path(__file__, 'test_cli_commands.py'), attribute='bar'
        ),
    )
    with server_process.create_ephemeral_client() as client:
        args = ['--grpc-host', client.host]
        if client.port:
            args.append('--grpc-port')
            args.append(client.port)
        if client.socket:
            args.append('--grpc-socket')
            args.append(client.socket)
        if pipeline_name:
            args.append('--pipeline')
            args.append(pipeline_name)

        yield args
    server_process.wait()


@contextmanager
def grpc_server_bar_pipeline_args():
    with default_instance() as instance:
        with grpc_server_bar_kwargs(pipeline_name='foo') as kwargs:
            yield kwargs, False, instance


# This iterates over a list of contextmanagers that can be used to contruct
# (cli_args, uses_legacy_repository_yaml_format, instance tuples)
def launch_command_contexts():
    for pipeline_target_args in valid_external_pipeline_target_args():
        yield args_with_default_instance(*pipeline_target_args)
        yield pytest.param(
            args_with_managed_grpc_instance(*pipeline_target_args), marks=pytest.mark.managed_grpc
        )
    yield pytest.param(grpc_server_bar_pipeline_args(), marks=pytest.mark.deployed_grpc)


def pipeline_python_origin_contexts():
    return [
        args_with_default_instance(pipeline_target_args)
        for pipeline_target_args in valid_pipeline_python_origin_target_args()
    ]


@contextmanager
def scheduler_instance(overrides=None):
    with seven.TemporaryDirectory() as temp_dir:
        with default_instance_tempdir(
            temp_dir,
            overrides=merge_dicts(
                {
                    'scheduler': {
                        'module': 'dagster.utils.test',
                        'class': 'FilesystemTestScheduler',
                        'config': {'base_dir': temp_dir},
                    }
                },
                overrides if overrides else {},
            ),
        ) as instance:
            yield instance


def managed_grpc_scheduler_instance():
    return scheduler_instance(overrides={"opt_in": {"local_servers": True}})


@contextmanager
def grpc_server_scheduler_cli_args():
    with scheduler_instance() as instance:
        with grpc_server_bar_cli_args() as args:
            yield args, instance


# Returns a list of contextmanagers that can be used to contruct
# (cli_args, instance) tuples for schedule calls
def schedule_command_contexts():
    return [
        args_with_instance(
            scheduler_instance(), ['-w', file_relative_path(__file__, 'workspace.yaml')]
        ),
        pytest.param(
            args_with_instance(
                managed_grpc_scheduler_instance(),
                ['-w', file_relative_path(__file__, 'workspace.yaml')],
            ),
            marks=pytest.mark.managed_grpc,
        ),
        pytest.param(grpc_server_scheduler_cli_args(), marks=pytest.mark.deployed_grpc),
    ]


# This iterates over a list of contextmanagers that can be used to contruct
# (cli_args, uses_legacy_repository_yaml_format, instance) tuples for backfill calls
def backfill_command_contexts():
    repo_args = {
        'noprompt': True,
        'workspace': (file_relative_path(__file__, 'repository_file.yaml'),),
    }
    return [
        args_with_instance(default_instance(), repo_args, True),
        pytest.param(
            args_with_instance(managed_grpc_instance(), repo_args, True),
            marks=pytest.mark.managed_grpc,
        ),
        pytest.param(grpc_server_backfill_args(), marks=pytest.mark.deployed_grpc),
    ]


@contextmanager
def grpc_server_backfill_args():
    with default_instance() as instance:
        with grpc_server_bar_kwargs() as args:
            yield merge_dicts(args, {'noprompt': True}), False, instance


def valid_pipeline_python_origin_target_args():
    return [
        {
            'workspace': None,
            'pipeline': 'foo',
            'python_file': file_relative_path(__file__, 'test_cli_commands.py'),
            'module_name': None,
            'attribute': 'bar',
        },
        {
            'workspace': None,
            'pipeline': 'foo',
            'python_file': file_relative_path(__file__, 'test_cli_commands.py'),
            'module_name': None,
            'attribute': 'bar',
            'working_directory': os.path.dirname(__file__),
        },
        {
            'workspace': None,
            'pipeline': 'foo',
            'python_file': None,
            'module_name': 'dagster_tests.cli_tests.command_tests.test_cli_commands',
            'attribute': 'bar',
        },
        {
            'workspace': None,
            'pipeline': None,
            'python_file': None,
            'module_name': 'dagster_tests.cli_tests.command_tests.test_cli_commands',
            'attribute': 'foo_pipeline',
        },
        {
            'workspace': None,
            'pipeline': None,
            'python_file': file_relative_path(__file__, 'test_cli_commands.py'),
            'module_name': None,
            'attribute': 'define_foo_pipeline',
        },
        {
            'workspace': None,
            'pipeline': None,
            'python_file': file_relative_path(__file__, 'test_cli_commands.py'),
            'module_name': None,
            'attribute': 'define_foo_pipeline',
            'working_directory': os.path.dirname(__file__),
        },
        {
            'workspace': None,
            'pipeline': None,
            'python_file': file_relative_path(__file__, 'test_cli_commands.py'),
            'module_name': None,
            'attribute': 'foo_pipeline',
        },
    ]


# [(cli_args, uses_legacy_repository_yaml_format)]
def valid_external_pipeline_target_args():
    return [
        (
            {
                'workspace': (file_relative_path(__file__, 'repository_file.yaml'),),
                'pipeline': 'foo',
                'python_file': None,
                'module_name': None,
                'attribute': None,
            },
            True,
        ),
        (
            {
                'workspace': (file_relative_path(__file__, 'repository_module.yaml'),),
                'pipeline': 'foo',
                'python_file': None,
                'module_name': None,
                'attribute': None,
            },
            True,
        ),
    ] + [(args, False) for args in valid_pipeline_python_origin_target_args()]


def valid_pipeline_python_origin_target_cli_args():
    return [
        ['-f', file_relative_path(__file__, 'test_cli_commands.py'), '-a', 'bar', '-p', 'foo'],
        [
            '-f',
            file_relative_path(__file__, 'test_cli_commands.py'),
            '-d',
            os.path.dirname(__file__),
            '-a',
            'bar',
            '-p',
            'foo',
        ],
        [
            '-m',
            'dagster_tests.cli_tests.command_tests.test_cli_commands',
            '-a',
            'bar',
            '-p',
            'foo',
        ],
        ['-m', 'dagster_tests.cli_tests.command_tests.test_cli_commands', '-a', 'foo_pipeline'],
        ['-f', file_relative_path(__file__, 'test_cli_commands.py'), '-a', 'define_foo_pipeline',],
        [
            '-f',
            file_relative_path(__file__, 'test_cli_commands.py'),
            '-d',
            os.path.dirname(__file__),
            '-a',
            'define_foo_pipeline',
        ],
    ]


# [(cli_args, uses_legacy_repository_yaml_format)]
def valid_external_pipeline_target_cli_args_no_preset():
    return [
        (['-w', file_relative_path(__file__, 'repository_file.yaml'), '-p', 'foo'], True),
        (['-w', file_relative_path(__file__, 'repository_module.yaml'), '-p', 'foo'], True),
        (['-w', file_relative_path(__file__, 'workspace.yaml'), '-p', 'foo'], False),
        (
            [
                '-w',
                file_relative_path(__file__, 'override.yaml'),
                '-w',
                file_relative_path(__file__, 'workspace.yaml'),
                '-p',
                'foo',
            ],
            False,
        ),
    ] + [(args, False) for args in valid_pipeline_python_origin_target_cli_args()]


def valid_external_pipeline_target_cli_args_with_preset():
    return valid_external_pipeline_target_cli_args_no_preset() + [
        (
            [
                '-f',
                file_relative_path(__file__, 'test_cli_commands.py'),
                '-d',
                os.path.dirname(__file__),
                '-a',
                'define_foo_pipeline',
                '--preset',
                'test',
            ],
            False,
        )
    ]


def test_run_list():
    runner = CliRunner()
    result = runner.invoke(run_list_command)
    assert result.exit_code == 0


def test_run_wipe_correct_delete_message():
    runner = CliRunner()
    result = runner.invoke(run_wipe_command, input="DELETE\n")
    assert 'Deleted all run history and event logs' in result.output
    assert result.exit_code == 0


def test_run_wipe_incorrect_delete_message():
    runner = CliRunner()
    result = runner.invoke(run_wipe_command, input="WRONG\n")
    assert 'Exiting without deleting all run history and event logs' in result.output
    assert result.exit_code == 0
