from pathlib import Path

import pytest
from cwl_utils.parser.cwl_v1_2 import (
    CommandLineTool,
    CommandOutputArraySchema,
    CommandOutputBinding,
    CommandOutputParameter,
    DockerRequirement,
    File,
    Workflow,
    WorkflowStep,
    StepInputExpressionRequirement,
    WorkflowStepInput,
    WorkflowOutputParameter,
)
from cwl_utils.parser import save

from CTADIRAC.Interfaces.Utilities.CWLTranslator import CWLTranslator
from CTADIRAC.Interfaces.Utilities.CWLUtilities import (
    LFN_PREFIX,
    LFN_DIRAC_PREFIX,
    LOCAL_PREFIX,
)

CVMFS_BASE_PATH = Path("/cvmfs/ctao.dpps.test")
DOCKER_PYTHON_TAG = "harbor/python:tag"


def test_init_class(mocker):
    cwl_workflow = "tests/resources/cwl/single_command_line_tool/process_dl0_dl1.cwl"
    cwl_inputs = (
        "tests/resources/cwl/single_command_line_tool/inputs_process_dl0_dl1.yaml"
    )

    mocker_load_document = mocker.patch(
        "CTADIRAC.Interfaces.Utilities.CWLTranslator.load_document"
    )
    mocker_load_inputfile = mocker.patch(
        "CTADIRAC.Interfaces.Utilities.CWLTranslator.load_inputfile"
    )

    CWLTranslator(cwl_workflow, cwl_inputs)
    mocker_load_document.assert_called_once()
    mocker_load_inputfile.assert_called_once()


def test_translate(mocker):
    # Check if translate method calls the right methods (CLT or Workflow)
    cwl_translator = CWLTranslator.__new__(CWLTranslator)

    # --------------------------------------------
    # Test CommandLineTool
    cwl_translator.transformed_cwl = CommandLineTool.__new__(CommandLineTool)
    mock_translate_clt = mocker.patch(
        "CTADIRAC.Interfaces.Utilities.CWLTranslator.CWLTranslator._translate_clt"
    )

    cwl_translator.translate(CVMFS_BASE_PATH, [])
    mock_translate_clt.assert_called_once_with(CVMFS_BASE_PATH, [])

    # --------------------------------------------
    # Test Workflow
    cwl_translator.transformed_cwl = Workflow.__new__(Workflow)
    mock_translate_workflow = mocker.patch(
        "CTADIRAC.Interfaces.Utilities.CWLTranslator.CWLTranslator._translate_workflow"
    )

    cwl_translator.translate(CVMFS_BASE_PATH, [])
    mock_translate_workflow.assert_called_once_with(CVMFS_BASE_PATH, [])


def test_translate_clt(mocker):
    cwl_translator = CWLTranslator.__new__(CWLTranslator)

    cwl_obj = CommandLineTool(
        inputs=[],
        outputs=[],
        hints=[DockerRequirement(dockerPull=DOCKER_PYTHON_TAG)],
        baseCommand="python",
    )
    cwl_translator.original_inputs = {}
    cwl_translator.transformed_cwl = cwl_obj
    cwl_translator.transformed_inputs = {}
    cwl_translator.output_data = []
    cwl_translator.output_sandbox = []
    cwl_translator.current_step_name = None
    cwl_translator.steps_outputs = {}
    mock_extract_output_files = mocker.patch(
        "CTADIRAC.Interfaces.Utilities.CWLTranslator.CWLTranslator._extract_output_files"
    )

    cwl_translator._translate_clt(CVMFS_BASE_PATH, [])
    assert cwl_translator.transformed_cwl == cwl_obj
    assert cwl_translator.transformed_inputs == {}
    assert cwl_translator.output_data == []
    assert cwl_translator.output_sandbox == []
    mock_extract_output_files.assert_called_once_with(
        cwl_translator.transformed_cwl, cwl_translator.transformed_inputs
    )


def test_translate_workflow():
    def create_workflow(requirement=False):
        return Workflow(
            steps=[
                WorkflowStep(
                    id="/some/path#step_1",
                    in_=[],
                    out=[],
                    run=CommandLineTool(
                        inputs=[],
                        outputs=[],
                        baseCommand="echo CLT1",
                        hints=[DockerRequirement(dockerPull=DOCKER_PYTHON_TAG)],
                    ),
                ),
                WorkflowStep(
                    id="/some/path#step_2",
                    in_=[WorkflowStepInput(id="/some/path#step_4", valueFrom="value")]
                    if requirement
                    else [],
                    out=[],
                    run=CommandLineTool(
                        inputs=[],
                        outputs=[],
                        baseCommand=["echo CLT2"],
                        hints=[DockerRequirement(dockerPull=DOCKER_PYTHON_TAG)],
                    ),
                ),
                WorkflowStep(
                    id="/some/path#step_3",
                    in_=[WorkflowStepInput(id="/some/path#step_5", valueFrom="value")]
                    if requirement
                    else [],
                    out=[],
                    run=CommandLineTool(
                        inputs=[], outputs=[], baseCommand=["echo CLT3"]
                    ),
                ),
            ],
            inputs=[],
            outputs=[],
            requirements=[StepInputExpressionRequirement()] if requirement else [],
        )

    def test_workflow(cwl_workflow):
        cwl_translator.transformed_cwl = cwl_workflow

        cwl_translator._translate_workflow(CVMFS_BASE_PATH, [])

        assert cwl_translator.transformed_cwl == cwl_workflow
        assert cwl_translator.transformed_inputs == {}
        assert cwl_translator.output_data == []
        assert cwl_translator.output_sandbox == []

    cwl_translator = CWLTranslator.__new__(CWLTranslator)

    cwl_translator.original_inputs = {}
    cwl_translator.transformed_inputs = {}
    cwl_translator.output_data = []
    cwl_translator.output_sandbox = []
    cwl_translator.current_step_name = None
    cwl_translator.steps_outputs = {}

    # No requirements workflow
    cwl_workflow = create_workflow()

    # Requirements workflow
    cwl_workflow_req = create_workflow(True)

    test_workflow(cwl_workflow)
    test_workflow(cwl_workflow_req)


@pytest.mark.parametrize(
    ("hints", "base_command", "expected_hints", "expected_base_command"),
    [
        (
            [DockerRequirement(dockerPull=DOCKER_PYTHON_TAG)],
            "python",
            [],
            ["apptainer", "run", str(CVMFS_BASE_PATH / DOCKER_PYTHON_TAG), "python"],
        )
    ],
)
def test_translate_docker_hints(
    hints, base_command, expected_hints, expected_base_command
):
    cwl_translator = CWLTranslator.__new__(CWLTranslator)
    cwl_obj = CommandLineTool(
        inputs=None, outputs=None, hints=hints, baseCommand=base_command
    )

    result = cwl_translator._translate_docker_hints(cwl_obj, CVMFS_BASE_PATH, [])

    assert result.hints == expected_hints
    assert result.baseCommand == expected_base_command


@pytest.mark.parametrize(
    ("input_cwl", "expected_results"),
    [
        (
            {"input1": File(path=LFN_PREFIX + "/ctao/test_lfn_file.txt")},
            {
                "transformed_inputs": {"input1": File(path="test_lfn_file.txt")},
                "input_sandbox": [],
                "input_data": [LFN_DIRAC_PREFIX + "/ctao/test_lfn_file.txt"],
            },
        ),
        (
            {"input1": File(path=LOCAL_PREFIX + "test_local_file.txt")},
            {
                "transformed_inputs": {"input1": File(path="test_local_file.txt")},
                "input_sandbox": [str(Path("test_local_file.txt").absolute())],
                "input_data": [],
            },
        ),
        (
            {
                "input1": [
                    File(path=LFN_PREFIX + "/ctao/test_lfn_file1.txt"),
                    File(path=LOCAL_PREFIX + "test_local_file1.txt"),
                ]
            },
            {
                "transformed_inputs": {
                    "input1": [
                        File(path="test_lfn_file1.txt"),
                        File(path="test_local_file1.txt"),
                    ]
                },
                "input_sandbox": [str(Path("test_local_file1.txt").absolute())],
                "input_data": [LFN_DIRAC_PREFIX + "/ctao/test_lfn_file1.txt"],
            },
        ),
        (
            {
                "input1": File(path=LFN_PREFIX + "/ctao/test_lfn_file2.txt"),
                "input2": File(path=LOCAL_PREFIX + "test_local_file2.txt"),
                "input3": [
                    File(path=LFN_PREFIX + "/ctao/test_lfn_file3.txt"),
                    File(path=LOCAL_PREFIX + "test_local_file3.txt"),
                ],
            },
            {
                "transformed_inputs": {
                    "input1": File(path="test_lfn_file2.txt"),
                    "input2": File(path="test_local_file2.txt"),
                    "input3": [
                        File(path="test_lfn_file3.txt"),
                        File(path="test_local_file3.txt"),
                    ],
                },
                "input_sandbox": [
                    str(Path("test_local_file2.txt").absolute()),
                    str(Path("test_local_file3.txt").absolute()),
                ],
                "input_data": [
                    LFN_DIRAC_PREFIX + "/ctao/test_lfn_file2.txt",
                    LFN_DIRAC_PREFIX + "/ctao/test_lfn_file3.txt",
                ],
            },
        ),
        (
            {
                "input1": [
                    File(path="some/path/test_local_file1.txt"),
                ]
            },
            {
                "transformed_inputs": {
                    "input1": [
                        File(path="test_local_file1.txt"),
                    ]
                },
                "input_sandbox": [
                    str(Path("some/path/test_local_file1.txt").absolute())
                ],
                "input_data": [],
            },
        ),
        (
            {
                "input1": File(path="some/path/test_local_file1.txt"),
            },
            {
                "transformed_inputs": {
                    "input1": File(path="test_local_file1.txt"),
                },
                "input_sandbox": [
                    str(Path("some/path/test_local_file1.txt").absolute())
                ],
                "input_data": [],
            },
        ),
        # Test that a input string containing "lfn" is correctly treated.
        (
            {
                "input1": File(path="some/path/test_local_file1.txt"),
                "input2": f"{LFN_PREFIX}/ctao/path/test.h5",
            },
            {
                "transformed_inputs": {
                    "input1": File(path="test_local_file1.txt"),
                    "input2": f"{LFN_PREFIX}/ctao/path/test.h5",
                },
                "input_sandbox": [
                    str(Path("some/path/test_local_file1.txt").absolute())
                ],
                "input_data": [],
            },
        ),
    ],
)
def test_extract_input_files(input_cwl, expected_results):
    cwl_translator = CWLTranslator.__new__(CWLTranslator)

    cwl_translator.input_data = []
    cwl_translator.input_sandbox = []
    cwl_translator.transformed_inputs = input_cwl

    cwl_translator._extract_input_files(input_cwl)

    assert save(cwl_translator.transformed_inputs) == save(
        expected_results["transformed_inputs"]
    )
    assert cwl_translator.input_sandbox == expected_results["input_sandbox"]
    assert cwl_translator.input_data == expected_results["input_data"]


@pytest.mark.parametrize(
    ("file_input", "expected_result", "expected_lfn"),
    [
        (
            File(path=LFN_PREFIX + "/ctao/test_lfn_file.txt"),
            LFN_DIRAC_PREFIX + "/ctao/test_lfn_file.txt",
            True,
        ),
        (
            File(path=LOCAL_PREFIX + "/home/user/test_local_file.txt"),
            "/home/user/test_local_file.txt",
            False,
        ),
        (
            LFN_PREFIX + "/ctao/test_lfn_str.txt",
            LFN_DIRAC_PREFIX + "/ctao/test_lfn_str.txt",
            True,
        ),
        (
            LOCAL_PREFIX + "/home/user/test_local_str.txt",
            "/home/user/test_local_str.txt",
            False,
        ),
        (File(), None, False),  # This will raise an exception
    ],
)
def test_translate_sandboxes_and_lfns(file_input, expected_result, expected_lfn):
    cwl_translator = CWLTranslator.__new__(CWLTranslator)

    if expected_result is None:
        with pytest.raises(KeyError, match="File path is not defined."):
            cwl_translator._translate_sandboxes_and_lfns(file_input)
    else:
        result, is_lfn = cwl_translator._translate_sandboxes_and_lfns(file_input)
        assert result == expected_result
        assert is_lfn == expected_lfn


@pytest.mark.parametrize(
    ("outputs", "expected_output_sandbox", "expected_output_data"),
    [
        (
            [
                CommandOutputParameter(
                    type_="File",
                    outputBinding=CommandOutputBinding(glob="/path/to/output1.txt"),
                )
            ],
            ["/path/to/output1.txt"],
            [],
        ),
        (
            [
                CommandOutputParameter(
                    type_="File",
                    outputBinding=CommandOutputBinding(
                        glob=LFN_PREFIX + "/path/to/output1.txt"
                    ),
                )
            ],
            [],
            [LFN_DIRAC_PREFIX + "/path/to/output1.txt"],
        ),
        (
            [
                CommandOutputParameter(
                    type_=CommandOutputArraySchema(type_="array", items=File),
                    outputBinding=CommandOutputBinding(
                        glob=[
                            LFN_PREFIX + "/path/to/output1.txt",
                            "/path/to/output2.txt",
                        ]
                    ),
                )
            ],
            ["/path/to/output2.txt"],
            [LFN_DIRAC_PREFIX + "/path/to/output1.txt"],
        ),
    ],
)
def test_extract_output_files(outputs, expected_output_sandbox, expected_output_data):
    cwl_translator = CWLTranslator.__new__(CWLTranslator)

    cwl_obj = CommandLineTool(inputs={}, outputs=outputs)
    cwl_translator.output_sandbox = []
    cwl_translator.output_data = []
    cwl_translator.current_step_name = None
    cwl_translator.steps_outputs = {}

    _ = cwl_translator._extract_output_files(cwl_obj, {})

    assert cwl_translator.output_sandbox == expected_output_sandbox
    assert cwl_translator.output_data == expected_output_data


STEP_INPUT_NAME = "input1"
INPUT_NAME_IN_INPUTS = "input_entry"


@pytest.mark.parametrize(
    ("step", "inputs", "expected_result"),
    [
        (
            WorkflowStep(
                id="/some/path#step_1",
                in_=[
                    WorkflowStepInput(
                        id=f"file:///path/to/#step_1/{STEP_INPUT_NAME}",
                        source=f"file:///path/to/#{INPUT_NAME_IN_INPUTS}",
                    )
                ],
                out=[],
                scatter=[STEP_INPUT_NAME],
                run=CommandLineTool(
                    inputs=[],
                    outputs=[],
                    baseCommand="echo CLT1",
                ),
            ),
            {
                INPUT_NAME_IN_INPUTS: {"class": "File", "path": "/path/to/input"},
                "other_input": "value",
            },
            {STEP_INPUT_NAME: {"class": "File", "path": "/path/to/input"}},
        )
    ],
)
def test_get_scattered_inputs(step, inputs, expected_result):
    cwl_translator = CWLTranslator.__new__(CWLTranslator)
    result = cwl_translator._get_scattered_inputs(step, inputs)
    assert result == expected_result


@pytest.mark.parametrize(
    ("step", "inputs", "expected_result"),
    [
        (
            WorkflowStep(
                id="/some/path#step_1",
                in_=[],
                out=[],
                run=CommandLineTool(
                    inputs=["input1"],
                    outputs=[
                        CommandOutputParameter(
                            id="/some/path#step_1/output1",
                            type_="File",
                            outputBinding=CommandOutputBinding(glob="$(inputs.input1)"),
                        ),
                        CommandOutputParameter(
                            id="/some/path#step_1/output2",
                            type_="string",
                            outputBinding=CommandOutputBinding(glob="static_glob"),
                        ),
                    ],
                    baseCommand="echo CLT1",
                ),
                scatter="input1",
            ),
            {
                "input1": [
                    {"class": "File", "path": "file1.txt"},
                    {"class": "File", "path": "file2.txt"},
                ],
            },
            {
                "step_1": {
                    "output1": [
                        {"class": "File", "path": "file1.txt"},
                        {"class": "File", "path": "file2.txt"},
                    ],
                    "output2": "static_glob",
                }
            },
        ),
        (
            WorkflowStep(
                id="/some/path#step_1",
                in_=[],
                out=[],
                run=CommandLineTool(
                    inputs=["input1"],
                    outputs=[
                        CommandOutputParameter(
                            id="/some/path#step_1/output1",
                            type_="File",
                            outputBinding=CommandOutputBinding(glob="$(inputs.input1)"),
                        ),
                        CommandOutputParameter(
                            id="/some/path#step_1/output2",
                            type_="File[]",
                            outputBinding=CommandOutputBinding(
                                glob=["file3.txt", "file4.txt"]
                            ),
                        ),
                    ],
                    baseCommand="echo CLT1",
                ),
                scatter="input1",
            ),
            {
                "input1": [
                    {"class": "File", "path": "file1.txt"},
                    {"class": "File", "path": "file2.txt"},
                ],
            },
            {
                "step_1": {
                    "output1": [
                        {"class": "File", "path": "file1.txt"},
                        {"class": "File", "path": "file2.txt"},
                    ],
                    "output2": [
                        {"class": "File", "path": "file3.txt"},
                        {"class": "File", "path": "file4.txt"},
                    ],
                }
            },
        ),
    ],
)
def test_evaluate_step_outputs(step, inputs, expected_result):
    cwl_translator = CWLTranslator.__new__(CWLTranslator)
    cwl_translator.current_step_name = "step_1"
    cwl_translator.evaluated_steps_outputs = {}
    cwl_translator.evaluated_steps_outputs.setdefault(
        cwl_translator.current_step_name, {}
    )
    cwl_translator._evaluate_step_outputs(step, inputs)
    assert cwl_translator.evaluated_steps_outputs == expected_result


@pytest.mark.parametrize(
    (
        "transformed_cwl",
        "evaluated_steps_outputs",
        "expected_output_sandbox",
        "expected_output_data",
    ),
    [
        (
            Workflow(
                inputs=[],
                outputs=[
                    WorkflowOutputParameter(
                        id="/some/path#output1",
                        outputSource="/some/path#step_1/output1",
                        type_="File",
                    ),
                    WorkflowOutputParameter(
                        id="/some/path#output2",
                        outputSource="/some/path#step_2/output2",
                        type_="File[]",
                    ),
                ],
                steps=[
                    WorkflowStep(
                        id="/some/path#step_1",
                        in_=[],
                        out=[],
                        run=CommandLineTool(
                            inputs=[],
                            outputs=[],
                            baseCommand="echo CLT1",
                            hints=[DockerRequirement(dockerPull=DOCKER_PYTHON_TAG)],
                        ),
                    )
                ],
            ),
            {
                "step_1": {
                    "output1": {"class": "File", "path": "/local/path/file1.txt"},
                },
                "step_2": {
                    "output2": [
                        {"class": "File", "path": f"{LFN_PREFIX}remote/path/file2.txt"},
                        {"class": "File", "path": "/local/path/file3.txt"},
                    ],
                },
            },
            ["/local/path/file1.txt", "/local/path/file3.txt"],
            [f"{LFN_DIRAC_PREFIX}remote/path/file2.txt"],
        ),
    ],
)
def test_extract_workflow_outputs(
    transformed_cwl,
    evaluated_steps_outputs,
    expected_output_sandbox,
    expected_output_data,
):
    # Initialize the class under test
    instance = CWLTranslator.__new__(CWLTranslator)

    # Set up the necessary state
    instance.transformed_cwl = transformed_cwl
    instance.evaluated_steps_outputs = evaluated_steps_outputs
    instance.output_sandbox = []
    instance.output_data = []

    # Call the method under test
    instance._extract_workflow_outputs()

    # Assert the results
    assert instance.output_sandbox == expected_output_sandbox
    assert instance.output_data == expected_output_data
