from pathlib import Path

import pytest

from lang_main import io
from lang_main.errors import (
    NoPerformableActionError,
    OutputInPipelineContainerError,
    WrongActionTypeError,
)
from lang_main.pipelines import base

PIPELINE_NAME = 'test'


@pytest.fixture(scope='module')
def working_dir() -> Path:
    work_dir = Path.cwd() / 'tests/work_dir'
    if not work_dir.exists():
        work_dir.mkdir()
    return work_dir


@pytest.fixture(scope='function')
def pipeline_container(working_dir) -> base.PipelineContainer:
    return base.PipelineContainer(name=PIPELINE_NAME, working_dir=working_dir)


@pytest.fixture(scope='function')
def pipeline(working_dir) -> base.Pipeline:
    return base.Pipeline(name=PIPELINE_NAME, working_dir=working_dir)


def test_empty_pipeline_container(pipeline_container, working_dir):
    container = pipeline_container
    assert container.name == PIPELINE_NAME
    assert container.working_dir == working_dir
    assert len(container.actions) == 0
    assert len(container.action_names) == 0
    assert len(container.action_skip) == 0
    assert container.curr_proc_idx == 1

    with pytest.raises(NoPerformableActionError):
        container.prep_run()

    assert container.post_run() is None


@pytest.mark.parametrize('skip', [True, False])
def test_pipeline_container_valid(pipeline_container, skip):
    test_string = 'test'

    def valid_action():  # pragma: no cover
        nonlocal test_string
        test_string += '_2'

    pipeline_container.add(valid_action, skip=skip)
    assert len(pipeline_container.actions) == 1
    assert len(pipeline_container.action_names) == 1
    assert len(pipeline_container.action_skip) == 1

    ret = pipeline_container.run()
    assert pipeline_container.curr_proc_idx == 2
    assert ret is None
    if skip:
        assert test_string == 'test'
    else:
        assert test_string == 'test_2'

    pipeline_container.prep_run()
    assert pipeline_container.curr_proc_idx == 1


def test_pipeline_container_invalid_action(pipeline_container):
    test_string = 'test'

    def invalid_action():
        nonlocal test_string
        test_string += '_2'
        new = 'ret'
        return new

    with pytest.raises(WrongActionTypeError):
        pipeline_container.add(test_string, skip=False)

    pipeline_container.add(invalid_action, skip=False)
    with pytest.raises(OutputInPipelineContainerError):
        pipeline_container.run()


def test_empty_pipeline(pipeline, working_dir):
    pipe = pipeline
    assert pipe.name == PIPELINE_NAME
    assert pipe.working_dir == working_dir
    assert len(pipe.actions) == 0
    assert len(pipe.action_names) == 0
    assert len(pipe.actions_kwargs) == 0
    assert len(pipe.save_results) == 0
    assert len(pipe.load_results) == 0
    assert pipe.curr_proc_idx == 1
    assert pipe._intermediate_result is None

    with pytest.raises(NoPerformableActionError):
        pipe.prep_run()

    assert pipe.post_run() is None


@pytest.mark.parametrize('alter_content', [True, False])
def test_pipeline_valid(pipeline, alter_content):
    pipe = pipeline
    test_string = 'test'

    # action preparation
    def valid_action(string, add_content=False):
        if add_content:
            string += '_2'
        return string

    pipe.add(valid_action, {'add_content': alter_content})
    assert len(pipe.actions) == 1
    assert len(pipe.action_names) == 1
    assert len(pipe.actions_kwargs) == 1
    assert len(pipe.action_skip) == 1
    assert len(pipe.save_results) == 1
    assert len(pipe.load_results) == 1
    assert pipe.save_results[0] == (False, None)
    assert pipe.load_results[0] == (False, None)
    # filenames and saving/loading
    target_filename = f'Pipe-{pipe.name}_Step-{pipe.curr_proc_idx}_valid_action'
    target_pth = (pipe.working_dir / target_filename).with_suffix('.pkl')
    ret_pth, action_name = pipe.get_result_path(0, filename=None)
    assert ret_pth == target_pth
    assert action_name == 'valid_action'
    filename = 'test'
    ret_pth, action_name = pipe.get_result_path(0, filename=filename)
    target_pth = (pipe.working_dir / filename).with_suffix('.pkl')
    assert ret_pth == target_pth
    assert action_name == 'valid_action'
    # load non-existing files
    with pytest.raises(FileNotFoundError):
        pipe.load_step(0, 'non_existing')

    # running
    ret = pipe.run(starting_values=(test_string,))
    assert isinstance(ret, tuple)
    assert pipe._intermediate_result == ret
    assert pipe.curr_proc_idx == 2
    assert ret is not None
    if alter_content:
        assert ret[0] == 'test_2'
    else:
        assert ret[0] == 'test'

    pipe.prep_run()
    assert pipe.curr_proc_idx == 1

    # load existing files
    loaded_res = pipe.load_step(0, None)
    assert loaded_res is not None
    assert isinstance(loaded_res, tuple)
    assert loaded_res[0] == ret[0]


def test_pipeline_valid_action_load(pipeline, working_dir):
    pipe = pipeline
    test_string = 'test'

    # action preparation
    def valid_action(string, add_content=False):  # pragma: no cover
        if add_content:
            string += '_2'
        return string

    pipe.add(valid_action, {'add_content': False}, load_result=True)
    assert len(pipe.actions) == 1
    assert len(pipe.action_names) == 1
    assert len(pipe.actions_kwargs) == 1
    assert len(pipe.action_skip) == 1
    assert len(pipe.save_results) == 1
    assert len(pipe.load_results) == 1
    assert pipe.save_results[0] == (False, None)
    assert pipe.load_results[0] == (True, None)

    ret = pipe.run(starting_values=(test_string,))
    assert isinstance(ret, tuple)
    assert pipe._intermediate_result == ret
    assert pipe.curr_proc_idx == 2
    assert ret is not None

    # load non-tuple result
    filename = 'non_tuple.pkl'
    save_pth = working_dir / filename
    io.save_pickle(test_string, save_pth)
    with pytest.raises(TypeError):
        pipe.load_step(0, filename)


def test_pipeline_multiple_actions(pipeline):
    pipe = pipeline
    test_string = 'test'

    # action preparation
    def valid_action(string, add_content=True):
        if add_content:
            string += '_2'
        return string

    def valid_action_2(string, add_content=True):
        if add_content:
            string += '_3'
        return string

    pipe.add(valid_action, {'add_content': True}, skip=True)
    pipe.add(valid_action, {'add_content': True})
    pipe.add(valid_action_2)
    assert len(pipe.actions) == 3
    assert len(pipe.action_names) == 3
    assert len(pipe.actions_kwargs) == 3
    assert len(pipe.action_skip) == 3
    assert len(pipe.save_results) == 3
    assert len(pipe.load_results) == 3
    assert pipe.save_results[1] == (False, None)
    assert pipe.load_results[1] == (False, None)

    ret = pipe.run(starting_values=(test_string,))
    assert isinstance(ret, tuple)
    assert pipe._intermediate_result == ret
    assert pipe.curr_proc_idx == 4
    assert ret is not None
    assert ret[0] == 'test_2_3'


def test_pipeline_invalid_action(pipeline):
    test_string = 'test'

    with pytest.raises(WrongActionTypeError):
        pipeline.add(test_string, skip=False)
