import filecmp
import warnings
from pathlib import Path
from unittest import mock

import pytest

with warnings.catch_warnings():
    warnings.simplefilter("ignore", UserWarning)
    try:
        import htcondor2 as htcondor
    except ImportError:
        import htcondor

from .. import DAG, Argument, File, Layer, Node, Option


def compare_submit_files(file1: Path, file2: Path) -> bool:
    """Compare two submit files by loading them as Submit objects.

    This provides a semantic comparison rather than exact string matching,
    making tests more robust to formatting changes.
    """
    # Read both files
    submit1_text = file1.read_text()
    submit2_text = file2.read_text()

    # Parse into Submit objects
    submit1 = htcondor.Submit(submit1_text)
    submit2 = htcondor.Submit(submit2_text)

    # Get all keys from both submits
    keys1 = set(submit1.keys())
    keys2 = set(submit2.keys())

    # Check if keys match
    if keys1 != keys2:
        missing_in_1 = keys2 - keys1
        missing_in_2 = keys1 - keys2
        if missing_in_1:
            print(f"Keys in {file2.name} but not in {file1.name}: {missing_in_1}")
        if missing_in_2:
            print(f"Keys in {file1.name} but not in {file2.name}: {missing_in_2}")
        return False

    # Compare values for each key
    for key in keys1:
        val1 = submit1.get(key, "").strip()
        val2 = submit2.get(key, "").strip()
        if val1 != val2:
            print(f"Value mismatch for key '{key}':")
            print(f"  {file1.name}: {val1}")
            print(f"  {file2.name}: {val2}")
            return False

    return True


@mock.patch("shutil.which", side_effect=lambda x: f"/path/to/{x}")
def test_dag_generation(mock_which, shared_datadir, tmp_path):
    # create DAG
    dag_name = "my_dag"
    dag = DAG(dag_name)

    # define job requirements
    requirements = {
        "environment": {
            "OMP_NUM_THREADS": 1,
        },
        "getenv": [
            "HOME",
            "USER",
        ],
        "request_cpus": 1,
        "request_memory": 2000,
        "requirements": [
            "HAS_CVMFS_oasis_opensciencegrid_org=TRUE",
        ],
        "transfer_plugins": {"custom": "custom_plugin.py"},
    }

    # create processing layer, add nodes
    process_layer = Layer("process_bins", submit_description=requirements)
    output_files = []
    for i in range(3):
        output_file = f"custom://output_{i}.txt"
        bin_idx = [3 * j + i for j in range(3)]
        node = Node(
            arguments=[
                Argument("job-index", i),
                Option("verbose"),
                Option("bins", bin_idx),
            ],
            inputs=Option("input", "data.txt"),
            outputs=Argument("output", output_file),
        )

        process_layer.append(node)
        output_files.append(output_file)

        # check command
        cmd = process_layer.command(node)
        bin_cmd = Option("bins", bin_idx).vars()
        expected = (
            f"process_bins {i} --verbose {bin_cmd} --input data.txt output_{i}.txt"
        )
        assert cmd == expected

    # add layer to DAG
    dag.attach(process_layer)

    # create combine layer, add node
    combine_layer = Layer("combine_bins", submit_description=requirements)
    combine_layer += Node(
        arguments=Option("verbose"),
        inputs=Argument("input", output_files),
        outputs=Argument("output", "combined.txt"),
    )

    # add layer to DAG
    dag.attach(combine_layer)

    # write DAG to disk
    dag.write(tmp_path, write_script=True)

    # write DAG graph to disk
    dag.visualize(tmp_path)

    # compare contents of generated files
    dag_filename = f"{dag_name}.dag"
    script_filename = f"{dag_name}.sh"
    graph_filename = f"{dag_name}.svg"

    # Compare DAG file
    assert filecmp.cmp(tmp_path / dag_filename, shared_datadir / dag_filename), (
        f"contents for {dag_filename} does not match expected output"
    )

    # Compare submit files semantically
    for job in ("process_bins", "combine_bins"):
        sub_filename = f"{job}.sub"
        assert compare_submit_files(
            tmp_path / sub_filename, shared_datadir / sub_filename
        ), f"submit description for {sub_filename} does not match expected output"

    # Compare script file (exact match)
    assert filecmp.cmp(tmp_path / script_filename, shared_datadir / script_filename), (
        f"contents for {script_filename} does not match expected output"
    )

    # Check graph file exists
    assert (tmp_path / graph_filename).exists()


@mock.patch("shutil.which", side_effect=lambda x: f"/path/to/{x}")
@pytest.mark.parametrize(
    ("name", "base_path", "transfer_files"),
    [
        ("rel", "relative/path/to", True),
        ("rel", "relative/path/to", False),
        ("abs", "/absolute/path/to", True),
        ("abs", "/absolute/path/to", False),
        ("url", "osdf:///path/to", True),
    ],
)
def test_dag_file_paths(
    mock_which, name, base_path, transfer_files, shared_datadir, tmp_path
):
    transfer_path = "transfer" if transfer_files else "no_transfer"
    expected_datadir = shared_datadir / name / transfer_path

    # create DAG
    dag_name = "count_dag"
    dag = DAG(dag_name)

    # define job requirements
    options = {
        "request_cpus": 1,
        "request_memory": 2000,
    }

    # create partitioning layer
    partition_layer = Layer(
        "partition_by_key", submit_description=options, transfer_files=transfer_files
    )
    output_files = []
    for i in range(3):
        input_files = [f"/abs/dir/data_{3 * i + j}.txt" for j in range(3)]
        output_file = f"{base_path}/output_{i}.txt"
        partition_layer += Node(
            arguments=["-j", 4],
            inputs=Argument("data", input_files),
            outputs=Argument("partitioned", output_file),
        )
        output_files.append(output_file)
    dag.attach(partition_layer)

    # create aggregation layer
    agg_layer = Layer(
        "aggregate_by_key", submit_description=options, transfer_files=transfer_files
    )
    for output_file in output_files:
        agg_layer += Node(
            arguments=Option("aggregate", "max"),
            inputs=Argument("data", output_file),
            outputs=File("aggregated-data", output_file),
        )
    dag.attach(agg_layer)

    # create count layer
    count_layer = Layer(
        "count_keys", submit_description=options, transfer_files=transfer_files
    )
    count_layer += Node(
        arguments=Option("verbose"),
        inputs=Argument("input", output_files),
        outputs=Argument("output", "combined.txt"),
    )
    dag.attach(count_layer)

    # write DAG to disk
    dag.write(tmp_path, write_script=True)

    # compare contents of generated files
    dag_filename = f"{dag_name}.dag"
    script_filename = f"{dag_name}.sh"

    # Compare DAG file
    assert filecmp.cmp(tmp_path / dag_filename, expected_datadir / dag_filename), (
        f"contents for {dag_filename} does not match expected output"
    )

    # Compare submit files semantically
    for job in ("partition_by_key", "aggregate_by_key", "count_keys"):
        sub_filename = f"{job}.sub"
        assert compare_submit_files(
            tmp_path / sub_filename, expected_datadir / sub_filename
        ), f"submit description for {sub_filename} does not match expected output"

    # Compare script file (exact match)
    assert filecmp.cmp(
        tmp_path / script_filename, expected_datadir / script_filename
    ), f"contents for {script_filename} does not match expected output"


@mock.patch("shutil.which", side_effect=lambda x: f"/path/to/{x}")
def test_duplicate_layer_names(mock_which, tmp_path):
    """Attaching layers with the same name appends a counter."""
    dag = DAG("dup_dag")

    for i in range(2):
        layer = Layer("same_name")
        layer += Node(
            arguments=Argument("idx", i),
            outputs=Argument("out", f"out_{i}.txt"),
        )
        dag.attach(layer)

    dag.write(tmp_path)

    # second layer should be renamed to same_name2
    assert "same_name" in dag._layers
    assert "same_name2" in dag._layers


@mock.patch("shutil.which", side_effect=lambda x: f"/path/to/{x}")
def test_multi_parent_edges(mock_which, tmp_path):
    """A child layer consuming outputs from two different parent layers."""
    dag = DAG("multi_parent")

    # parent layer A
    layer_a = Layer("parent_a")
    layer_a += Node(outputs=Argument("out", "a_out.txt"))
    dag.attach(layer_a)

    # parent layer B
    layer_b = Layer("parent_b")
    layer_b += Node(outputs=Argument("out", "b_out.txt"))
    dag.attach(layer_b)

    # child depends on both
    child = Layer("child")
    child += Node(
        inputs=Argument("data", ["a_out.txt", "b_out.txt"]),
        outputs=Argument("out", "child_out.txt"),
    )
    dag.attach(child)

    dag.write(tmp_path)

    # child should have edges from both parents
    assert "parent_a" in dag._all_edges["child"]
    assert "parent_b" in dag._all_edges["child"]


@mock.patch("shutil.which", side_effect=lambda x: f"/path/to/{x}")
def test_visualize(mock_which, tmp_path):
    """Test DAG visualization produces an output file."""
    pytest.importorskip("graphviz")

    dag = DAG("viz_dag")

    process_layer = Layer("process")
    for i in range(2):
        process_layer += Node(
            inputs=Option("input", "data.txt"),
            outputs=Argument("out", f"out_{i}.txt"),
        )
    dag.attach(process_layer)

    combine_layer = Layer("combine")
    combine_layer += Node(
        inputs=Argument("data", ["out_0.txt", "out_1.txt"]),
    )
    dag.attach(combine_layer)

    dag.write(tmp_path)
    dag.visualize(tmp_path)

    assert (tmp_path / "viz_dag.svg").exists()


@mock.patch("shutil.which", side_effect=lambda x: f"/path/to/{x}")
def test_absolute_log_dir(mock_which, tmp_path):
    """Layers with absolute log_dir create that directory directly."""
    dag = DAG("abs_log")
    abs_log = str(tmp_path / "absolute_logs")

    layer = Layer("my_exec", log_dir=abs_log)
    layer += Node(arguments=Argument("x", 1))
    dag.attach(layer)

    dag.write(tmp_path)

    assert Path(abs_log).is_dir()
