# Some Useful Data Structures -------------------------------------------------

# Yes, there are four types of caves. Large caves, small caves, start caves, and
# end caves. Yes, they have different fields. Sue me. Each cave indicates the size, 
# name, and index of the cave, if it's important for that type of cave. The 
# indices are used later to determine whether a cave has already been visited.
abstract type Cave end

struct StartCave <: Cave end
struct EndCave   <: Cave end
struct LargeCave <: Cave name::String end
struct SmallCave <: Cave
    name::String
    index::Int
end

# Given a name and index, produce the appropriate type of Cave depending
# on whether the name is uppercase or not.
issmallcave(S) = all(islowercase, S) && S ∉ ["start", "end"]
Cave(s::AbstractString, idx::Int)::SmallCave = SmallCave(s, idx)
function Cave(s::AbstractString, ::Nothing)::Cave
    s == "start" && return StartCave()
    s == "end"   && return EndCave()
    return LargeCave(s)
end


# Ingest the Data File ---------------------------------------------------------

# Read in each line of the input file and parse it as an 'edge'. The final product
# here will be a Dict, where the keys are `Cave`s and the values are the list of
# other `Cave`s the key Cave is connected to. This behaves similarly to an 
# adjacency list, except you can look up cave connections in O(1) time.
function ingest(path)
    paths = open(path) do f
        [split(line, "-") for line in readlines(f)]
    end

    cavemap = Dict{Cave, Vector{Cave}}()
    indexes = Dict{String,Int}()
    for path in paths
        # Need to add entries to the output dictionary for both paths, from
        # the left cave to the right, and from the right cave to the left.
        (left, right) = path

        # Get the appropriate index for each cave. If there's already an 
        # assigned index for that cave name, just use it, otherwise use one more
        # than the total number of caves already indexed. Only do this if the
        # cave name is all lowercase, a small cave.
        if issmallcave(left) && get(indexes, left, 0) == 0
            indexes[left] = length(indexes) + 1
        end
        if issmallcave(right) && get(indexes, right, 0) == 0
            indexes[right] = length(indexes) + 1
        end

        # Make the `Caves`!
        leftcave  = Cave(left,  get(indexes, left, nothing))
        rightcave = Cave(right, get(indexes, right, nothing))

        # Add the left cave to the cavemap
        destinations = get(cavemap, leftcave, Vector{Cave}())
        push!(destinations, rightcave)
        cavemap[leftcave] = destinations

        # Add the right cave to the cavemap
        destinations = get(cavemap, rightcave, Vector{Cave}())
        push!(destinations, leftcave)
        cavemap[rightcave] = destinations
    end
    return cavemap
end

# Some Useful Helper Functions and Types ---------------------------------------

# Just some handy type aliases for clarity
const CaveMap = Dict{Cave, Vector{Cave}}
const ExploreStack = Vector{Tuple{BitVector,Cave}}


# We only need to keep track of the small caves. If moving into a large
# cave, we only need to pass back the set of visited Caves. If we're moving
# into the end cave, then we don't need to know about the visited caves 
# anymore, we can just use an empty BitVector.
nextpath(visited::BitVector, ::LargeCave) = visited
nextpath(::BitVector, ::EndCave) = BitVector()
function nextpath(visited::BitVector, nextcave::SmallCave) 
    visited = deepcopy(visited)
    visited[nextcave.index] = true
    return visited
end

# Determines whether we should skip entering a cave. We should always enter a 
# Large cave (don't skip it) or the end cave, never re-enter the start cave,
# and only enter a small cave if we've never been inside it before (meaning
# it's corresponding index in the visited vector will be false).
shouldskip(::BitVector, ::LargeCave) = false
shouldskip(::BitVector, ::EndCave) = false
shouldskip(::BitVector, ::StartCave) = true
shouldskip(visited::BitVector, cave::SmallCave) = visited[cave.index]


# Solve Part One ---------------------------------------------------------------

# This is a pretty standard implementation of a depth-first search, with a bit
# of a twist. Most of the twist is handled by the different helper functions
# and types implemented above.
function part1(input)
    # Start by creating a "stack" and initializing it with the "start" cave 
    # and a BitVector large enough to hold the indices of all the small caves
    smallcavecount = mapreduce(C -> C isa SmallCave, +, keys(input))
    stack = ExploreStack([(falses(smallcavecount), StartCave())])
    paths = 0

    # While there are still caves to explore...
    while !isempty(stack)
        # Pop the last cave and the list of visited caves off the stack...
        (visitedsofar, cave) = pop!(stack)

        # If we've reached the end, just add one to our count of paths
        # and move on to the next item in the stack
        if cave isa EndCave
            paths += 1
            continue
        end
        
        # Otherwise, get the list of all the caves the current cave is
        # connected to, and add them to the stack to be visited next
        for nextcave in get(input, cave, [])
            # If we found the start, or if our next cave is in our list of
            # visited caves, don't go there.
            shouldskip(visitedsofar, nextcave) && continue

            # We only need to keep track of the small caves we've visited
            visited = nextpath(visitedsofar, nextcave)
            push!(stack, (visited, nextcave))
        end
    end
    return paths
end

# Some More Useful Structs and Methods -----------------------------------------

# We need to change the types of our list of visited caves and exploration 
# stack to take advantage of the new methods for identifying the next path
# and whether we should skip the next cave.
const Path = Tuple{BitVector,Bool}
const ExploreStackTwo = Vector{Tuple{Path,Cave}}

# The logic is the same for large caves and end caves, but now we need to
# account for whether we've visited a single small cave twice when 
# updating our `Path`
nextpath(path::Path, ::LargeCave)= path
nextpath(::Path, ::EndCave) = (BitVector(), false)
function nextpath(path::Path, nextcave::SmallCave)
    (visited, twovisits) = path
    twovisits = twovisits || visited[nextcave.index]
    visited = deepcopy(visited)
    visited[nextcave.index] = true
    return (visited, twovisits)
end

# The logic for large caves, end caves, and start caves remains unchanged, but
# we needed to define them again because our first version was too specialized.
# These versions using `::Any` would work for part one and part two. The big
# change is the logic for determining whether to skip a small cave.
shouldskip(::Any, ::LargeCave) = false
shouldskip(::Any, ::EndCave) = false
shouldskip(::Any, ::StartCave) = true
function shouldskip(path::Path, cave::SmallCave) 
    (visited, twovisits) = path
    return visited[cave.index] && twovisits
end


# Solve Part Two ---------------------------------------------------------------

# This bit is *exactly* the same as part one. Well, except for the first few
# lines. Here, we've changed the types of the arguments in our stack to use
# the new logic methods from above. Otherwise, the algorithm is exactly the
# same.
function part2(input)
    # Start by creating a "stack" and initializing it with the "start" cave , 
    # a BitVector large enough to hold the indices of all the small caves, and
    # a boolean value to indicate whether we've seen a small cave twice.
    smallcavecount = mapreduce(C -> C isa SmallCave, +, keys(input))
    initialpath = (falses(smallcavecount), false)
    stack = ExploreStackTwo([(initialpath, StartCave())])
    paths = 0

    # While there are still caves to explore...
    while !isempty(stack)
        # Pop the last cave and the `Path` off the stack...
        (pathsofar, cave) = pop!(stack)

        # If we've reached the end, just add one to our count of paths
        # and move on to the next item in the stack
        if cave isa EndCave
            paths += 1
            continue
        end
        
        # Otherwise, get the list of all the caves the current cave is
        # connected to, and add them to the stack to be visited next
        for nextcave in get(input, cave, [])
            # If we found the start, or if our puzzle logic tells us to skip
            # the next cave, don't go there.
            shouldskip(pathsofar, nextcave) && continue

            # We only need to keep track of the small caves we've visited, still
            path = nextpath(pathsofar, nextcave)
            push!(stack, (path, nextcave))
        end
    end
    return paths
end