# for graph mutating functions we need to do a few things
# 1: mute the graph (duh)
# 2: keep track of what was changed for the diff (if track == true)
# 3: invalidate operation caches

"""
    insert_node!(graph::DAG, node::Node; track = true, invalidate_cache = true)

Insert the node into the graph.

## Keyword Arguments
`track::Bool`: Whether to add the changes to the [`DAG`](@ref)'s [`Diff`](@ref). Should be set `false` in parsing or graph creation functions for performance.

`invalidate_cache::Bool`: Whether to invalidate caches associated with the changes. Should also be turned off for graph creation or parsing.

See also: [`remove_node!`](@ref), [`insert_edge!`](@ref), [`remove_edge!`](@ref)
"""
function insert_node!(graph::DAG, node::Node; track = true, invalidate_cache = true)
    # 1: mute
    push!(graph.nodes, node)

    # 2: keep track
    if (track)
        push!(graph.diff.addedNodes, node)
    end

    # 3: invalidate caches
    if (!invalidate_cache)
        return node
    end
    push!(graph.dirtyNodes, node)

    return node
end

"""
    insert_edge!(graph::DAG, node1::Node, node2::Node; track = true, invalidate_cache = true)

Insert the edge between node1 (child) and node2 (parent) into the graph.

## Keyword Arguments
`track::Bool`: Whether to add the changes to the [`DAG`](@ref)'s [`Diff`](@ref). Should be set `false` in parsing or graph creation functions for performance.

`invalidate_cache::Bool`: Whether to invalidate caches associated with the changes. Should also be turned off for graph creation or parsing.

See also: [`insert_node!`](@ref), [`remove_node!`](@ref), [`remove_edge!`](@ref)
"""
function insert_edge!(graph::DAG, node1::Node, node2::Node; track = true, invalidate_cache = true)
    #@assert (node2 ∉ parents(node1)) && (node1 ∉ children(node2)) "Edge to insert already exists"

    # 1: mute
    # edge points from child to parent
    push!(node1.parents, node2)
    push!(node2.children, node1)

    # 2: keep track
    if (track)
        push!(graph.diff.addedEdges, make_edge(node1, node2))
    end

    # 3: invalidate caches
    if (!invalidate_cache)
        return nothing
    end

    invalidate_operation_caches!(graph, node1)
    invalidate_operation_caches!(graph, node2)

    push!(graph.dirtyNodes, node1)
    push!(graph.dirtyNodes, node2)

    return nothing
end

"""
    remove_node!(graph::DAG, node::Node; track = true, invalidate_cache = true)

Remove the node from the graph.

## Keyword Arguments
`track::Bool`: Whether to add the changes to the [`DAG`](@ref)'s [`Diff`](@ref). Should be set `false` in parsing or graph creation functions for performance.

`invalidate_cache::Bool`: Whether to invalidate caches associated with the changes. Should also be turned off for graph creation or parsing.

See also: [`insert_node!`](@ref), [`insert_edge!`](@ref), [`remove_edge!`](@ref)
"""
function remove_node!(graph::DAG, node::Node; track = true, invalidate_cache = true)
    #@assert node in graph.nodes "Trying to remove a node that's not in the graph"

    # 1: mute
    delete!(graph.nodes, node)

    # 2: keep track
    if (track)
        push!(graph.diff.removedNodes, node)
    end

    # 3: invalidate caches
    if (!invalidate_cache)
        return nothing
    end

    invalidate_operation_caches!(graph, node)
    delete!(graph.dirtyNodes, node)

    return nothing
end

"""
    remove_edge!(graph::DAG, node1::Node, node2::Node; track = true, invalidate_cache = true)

Remove the edge between node1 (child) and node2 (parent) into the graph.

## Keyword Arguments
`track::Bool`: Whether to add the changes to the [`DAG`](@ref)'s [`Diff`](@ref). Should be set `false` in parsing or graph creation functions for performance.

`invalidate_cache::Bool`: Whether to invalidate caches associated with the changes. Should also be turned off for graph creation or parsing.

See also: [`insert_node!`](@ref), [`remove_node!`](@ref), [`insert_edge!`](@ref)
"""
function remove_edge!(graph::DAG, node1::Node, node2::Node; track = true, invalidate_cache = true)
    # 1: mute
    pre_length1 = length(node1.parents)
    pre_length2 = length(node2.children)

    #TODO: filter is very slow
    for i in eachindex(node1.parents)
        if (node1.parents[i] == node2)
            splice!(node1.parents, i)
            break
        end
    end

    for i in eachindex(node2.children)
        if (node2.children[i] == node1)
            splice!(node2.children, i)
            break
        end
    end

    #=@assert begin
        removed = pre_length1 - length(node1.parents)
        removed <= 1
    end "removed more than one node from node1's parents"=#

    #=@assert begin
        removed = pre_length2 - length(children(node2))
        removed <= 1
    end "removed more than one node from node2's children"=#

    # 2: keep track
    if (track)
        push!(graph.diff.removedEdges, make_edge(node1, node2))
    end

    # 3: invalidate caches
    if (!invalidate_cache)
        return nothing
    end

    invalidate_operation_caches!(graph, node1)
    invalidate_operation_caches!(graph, node2)
    if (node1 in graph)
        push!(graph.dirtyNodes, node1)
    end
    if (node2 in graph)
        push!(graph.dirtyNodes, node2)
    end

    return nothing
end

function replace_children!(task::FusedComputeTask, before, after)
    replacedIn1 = length(findall(x -> x == before, task.t1_inputs))
    replacedIn2 = length(findall(x -> x == before, task.t2_inputs))

    #@assert replacedIn1 >= 1 || replacedIn2 >= 1 "Nothing to replace while replacing $before with $after in $(task.t1_inputs...) and $(task.t2_inputs...)"

    replace!(task.t1_inputs, before => after)
    replace!(task.t2_inputs, before => after)

    # recursively descend down the tree, but only in the tasks where we're replacing things
    if replacedIn1 > 0
        replace_children!(task.first_task, before, after)
    end
    if replacedIn2 > 0
        replace_children!(task.second_task, before, after)
    end

    return nothing
end

function replace_children!(task::AbstractTask, before, after)
    return nothing
end

function update_child!(graph::DAG, n::Node, child_before::Symbol, child_after::Symbol; track = true)
    # only need to update fused compute tasks
    if !(typeof(task(n)) <: FusedComputeTask)
        return nothing
    end

    taskBefore = copy(task(n))

    #=if !((child_before in task(n).t1_inputs) || (child_before in task(n).t2_inputs))
        println("------------------ Nothing to replace!! ------------------")
        child_ids = Vector{String}()
        for child in children(n)
            push!(child_ids, "$(child.id)")
        end
        println("From $(child_before) to $(child_after) in $n with children $(child_ids)")
        @assert false
    end=#

    replace_children!(task(n), child_before, child_after)

    #=if !((child_after in task(n).t1_inputs) || (child_after in task(n).t2_inputs))
        println("------------------ Did not replace anything!! ------------------")
        child_ids = Vector{String}()
        for child in children(n)
            push!(child_ids, "$(child.id)")
        end
        println("From $(child_before) to $(child_after) in $n with children $(child_ids)")
        @assert false
    end=#

    # keep track
    if (track)
        push!(graph.diff.updatedChildren, (n, taskBefore))
    end
end

"""
    get_snapshot_diff(graph::DAG)

Return the graph's [`Diff`](@ref) since last time this function was called.

See also: [`revert_diff!`](@ref), [`AppliedOperation`](@ref) and [`revert_operation!`](@ref)
"""
function get_snapshot_diff(graph::DAG)
    return swapfield!(graph, :diff, Diff())
end

"""
    invalidate_caches!(graph::DAG, operation::NodeFusion)

Invalidate the operation caches for a given [`NodeFusion`](@ref).

This deletes the operation from the graph's possible operations and from the involved nodes' own operation caches.
"""
function invalidate_caches!(graph::DAG, operation::NodeFusion)
    delete!(graph.possibleOperations, operation)

    # delete the operation from all caches of nodes involved in the operation
    # TODO: filter is very slow
    for n in [1, 3]
        for i in eachindex(operation.input[n].nodeFusions)
            if operation == operation.input[n].nodeFusions[i]
                splice!(operation.input[n].nodeFusions, i)
                break
            end
        end
    end

    operation.input[2].nodeFusion = missing

    return nothing
end

"""
    invalidate_caches!(graph::DAG, operation::NodeReduction)

Invalidate the operation caches for a given [`NodeReduction`](@ref).

This deletes the operation from the graph's possible operations and from the involved nodes' own operation caches.
"""
function invalidate_caches!(graph::DAG, operation::NodeReduction)
    delete!(graph.possibleOperations, operation)

    for node in operation.input
        node.nodeReduction = missing
    end

    return nothing
end

"""
    invalidate_caches!(graph::DAG, operation::NodeSplit)

Invalidate the operation caches for a given [`NodeSplit`](@ref).

This deletes the operation from the graph's possible operations and from the involved nodes' own operation caches.
"""
function invalidate_caches!(graph::DAG, operation::NodeSplit)
    delete!(graph.possibleOperations, operation)

    # delete the operation from all caches of nodes involved in the operation
    # for node split there is only one node
    operation.input.nodeSplit = missing

    return nothing
end

"""
    invalidate_operation_caches!(graph::DAG, node::ComputeTaskNode)

Invalidate the operation caches of the given node through calls to the respective [`invalidate_caches!`](@ref) functions.
"""
function invalidate_operation_caches!(graph::DAG, node::ComputeTaskNode)
    if !ismissing(node.nodeReduction)
        invalidate_caches!(graph, node.nodeReduction)
    end
    if !ismissing(node.nodeSplit)
        invalidate_caches!(graph, node.nodeSplit)
    end
    while !isempty(node.nodeFusions)
        invalidate_caches!(graph, pop!(node.nodeFusions))
    end
    return nothing
end

"""
    invalidate_operation_caches!(graph::DAG, node::DataTaskNode)

Invalidate the operation caches of the given node through calls to the respective [`invalidate_caches!`](@ref) functions.
"""
function invalidate_operation_caches!(graph::DAG, node::DataTaskNode)
    if !ismissing(node.nodeReduction)
        invalidate_caches!(graph, node.nodeReduction)
    end
    if !ismissing(node.nodeSplit)
        invalidate_caches!(graph, node.nodeSplit)
    end
    if !ismissing(node.nodeFusion)
        invalidate_caches!(graph, node.nodeFusion)
    end
    return nothing
end
