using ParserCombinator

const PValveList = (E"s " + Repeat(p"[A-Z]{2}" + Repeat(E", ", 0, 1)) + Eos()) | (E" " + p"[A-Z]{2}" + Eos());
const PValveSpec = E"Valve " + p"[A-Z]{2}" + E" has flow rate=" + PInt64() + P"; tunnels? leads? to valve" + PValveList;

struct ValveSpec
    name::String
    flowrate::Int
    tunnels::Vector{String}
end

function parse_valve_line(s::String)::ValveSpec
    vs = parse_one(s, PValveSpec);
    ValveSpec(String(vs[1]), vs[2]::Int, [String(s) for s in vs[3:end]])
end

struct Peer
    valve::ValveSpec
    dist::Int
end

struct Node
    valve::ValveSpec
    routes::Dict{String, Peer}
end

function build_graph(lines::IOStream)::Dict{String,Node}
    l = [Node(parse_valve_line(s), Dict()) for s in eachline(lines)];
    Dict(n.valve.name => n for n in l)
end

function update_connections(neigh::Node, dist::Int, mine::Dict{String, Peer})::Int
    count = 0;
    for (s, p) = mine
        if s == neigh.valve.name
            continue
        end
        if !(s in keys(neigh.routes)) || neigh.routes[s].dist > p.dist+dist
            pp = Peer(p.valve, p.dist+dist);
            neigh.routes[s] = pp;
            count += 1;
        end
    end
    count
end

function populate_graph(g::Dict{String,Node})
    for (id, valve) in g
        for neigh in valve.valve.tunnels
            # Insert neighbors with distance = 1.
            valve.routes[neigh] = Peer(g[neigh].valve, 1);
        end
    end

    count = 1;
    # While information propagates through the graph, keep updating.
    while count > 0
        count = 0;
        for (id, node) = g
            for (id2, node2) = node.routes
                if id == id2
                    continue
                end
                count += update_connections(g[id2], node2.dist, node.routes);
            end
        end
        @show count
    end
end

# Sort key
"""Return the total flow expected for node `node`, assuming it is opened at `time`, and the best path from it
is followed (this works recursively).
"""
function expected_total_flow(time::Int, node::Node, g::Dict{String,Node}, opened::Vector{String}=[]; level=-1)::Tuple{String,Int,Int}
    if level == 0
        return ("", -1, -1);
    end
    if in(node.valve.name, opened)
        return ("", -1, -1);
    end
    self = (30 - time - 1) * node.valve.flowrate;
    push!(opened, node.valve.name);
    opts = [(vn, p.dist, expected_total_flow(time+1+p.dist, g[vn], g, opened; level=level-1)[3]) for (vn, p) in node.routes if p.valve.flowrate > 0];
    if level == 7
        @show opts
    end
    pop!(opened);
    _, bestix = findmax(t -> t[3], opts);
    (opts[bestix][1], opts[bestix][2], opts[bestix][3]+self)
end

"""
Starting at AA, check which node is best by exploring all possible paths up to a certain depth,
visiting nodes by preference.
"""
function trace_best_path(g::Dict{String,Node}; init="AA")
    opened = String[];
    active_pressure_relief = 0;
    pressure_relieved = 0;
    time = 0;
    last = init;
    while length(g) > length(opened)
        @show (time, last, opened)
        next, dist, score = expected_total_flow(time, g[last], g, opened; level=7);
        @show (next, dist, score);

        if in(next, opened)
            break;
        end

        pressure_relieved += (dist+1) * active_pressure_relief;
        println("Relieving $(dist+1) x $active_pressure_relief...");
        active_pressure_relief += g[next].valve.flowrate;

        push!(opened, last);
        last = next;
        time += dist + 1;
    end
    pressure_relieved += (30 - time) * active_pressure_relief;
    println("Relieving $(30-time) x $active_pressure_relief...");
    println("Total relief: $pressure_relieved");
end

function print_graph(g::Dict{String,Node})
    for (id, n) = g
        println("=== $id === ");
        println(" flowrate: $(n.valve.flowrate)");
        println(" neighbors:");
        for (id2, neigh) = n.routes
            println("   $id2 => $(neigh.dist)");
        end
    end
end

open("16/test_input.txt"; read=true) do fh
    g = build_graph(fh);
    populate_graph(g);
    @time trace_best_path(g);
end

open("16/input.txt"; read=true) do fh
    g = build_graph(fh);
    @time populate_graph(g);
    @time trace_best_path(g);
end