import Base: ==, checkbounds, collect, eltype, filter, first, firstindex, getindex,
    in, isempty, iterate, keys, last, lastindex, length, map, pairs, show,
    IteratorEltype, IteratorSize

export NamingContext,
    freenames,
    freshnames,
    pushfirst,
    split


"""
    UniqueIndices(start, log_skip)

Immutable sequence of indices which can be kept linear in a context by splitting:

````
[1, 2, 3, ...]
-> ([1, 3, 5, ...], [2, 4, 6, ...])
-> (([1, 5, 9, ...], [3, 7, 11, ...]), [2, 4, 6, ...])
-> etc.
````
"""
struct UniqueIndices
    start::Int
    log_skip::Int
end

nextindex(u::UniqueIndices) = UniqueIndices(u.start + 2^u.log_skip, u.log_skip)

function splitindex(u::UniqueIndices)
    left = UniqueIndices(u.start, u.log_skip + 1)
    right = UniqueIndices(u.start + 2^u.log_skip, u.log_skip + 1)
    (left, right)
end


"""
    NamingContext([freenames]; namehint = :x)

Keeps track of free variables and their indices, and can generate fresh names based on them.
Fresh names are generated based on `namehint` and different from all `freenames`.

This behaves mostly like an immutable array of `Name`s.
"""
struct NamingContext
    freenames::Vector{Name}
    namehint::Name
    indices::UniqueIndices
end

NamingContext(freenames = Name[]; namehint = :x) =
    NamingContext(reverse!(collect(freenames)), namehint, UniqueIndices(1, 0))

freenames(Γ::NamingContext) = Γ.freenames

==(Γ₁::NamingContext, Γ₂::NamingContext) = Γ₁.freenames == Γ₂.freenames
checkbounds(::Type{Bool}, Γ::NamingContext, i::Int) = firstindex(Γ) ≤ i ≤ lastindex(Γ)
# cat(Γ₁::NamingContext, Γs::NamingContext...) =
#     NamingContext(cat(Γ₁.freenames, freenames.(Γs)...))
collect(Γ::NamingContext) = Γ.freenames
# copy(Γ::NamingContext) = NamingContext(copy(Γ.freenames))
eltype(::Type{NamingContext}) = Symbol
first(Γ::NamingContext) = firstindex(Γ.freenames)
firstindex(Γ::NamingContext) = firstindex(Γ.freenames)
filter(f, Γ::NamingContext) = NamingContext(filter(f, Γ.freenames))
getindex(Γ::NamingContext, key) = getindex(Γ.freenames, reverseind(Γ.freenames, key))
in(name, Γ::NamingContext) = name in Γ.freenames
isempty(Γ::NamingContext) = isempty(Γ.freenames)
keys(Γ::NamingContext) = keys(Γ.freenames)
last(Γ::NamingContext) = last(Γ.freenames)
lastindex(Γ::NamingContext) = lastindex(Γ.freenames)
length(Γ::NamingContext) = length(Γ.freenames)
map(f, Γs::NamingContext...) = NamingContext(map(f, freenames.(Γs)...))
pairs(Γ::NamingContext) = zip(keys(Γ.freenames), Iterators.reverse(Γ.freenames))

function show(io::IO, Γ::NamingContext)
    print(io, "NamingContext([")
    join(io, reverse(sprint.(show, Γ.freenames)), ", ")
    print(io, "])")
end


struct FreshNames
    Γ::NamingContext
end

iterate(iter::FreshNames) = iterate(iter, iter.Γ.indices)
function iterate(iter::FreshNames, state)
    while true
        newname = Symbol(iter.Γ.namehint, state.start)
        if newname ∉ iter.Γ.freenames
            return newname, nextindex(state)
        end
        state = nextindex(state)
    end
end

IteratorSize(::Type{FreshNames}) = Base.IsInfinite()
IteratorEltype(::Type{FreshNames}) = Base.HasEltype()
eltype(::Type{FreshNames}) = Symbol


"""
    freshnames(Γ)

An (infinite) iterator for the fresh names that can be generated from `Γ`.
"""
freshnames(Γ::NamingContext) = FreshNames(Γ)


"""
    pushfirst(Γ, names...)

Construct a new `NamingContext` by adding `names` as free names to `Γ`.  The new names will be added 
to the beginning, thus having the lowest index in the new context.
"""
pushfirst(Γ::NamingContext, names...) =
    NamingContext([Γ.freenames; names...], Γ.namehint, Γ.indices)


"""
    freshname(Γ)

Generate a new name fresh in `Γ`, and a new context `Γ′`, which is `Γ` updated by this name.
Names are constructed based on `Γ.namehint`.
"""
function freshname(Γ::NamingContext)
    (newname, newindex) = iterate(FreshNames(Γ))
    Γ′ = NamingContext([Γ.freenames; newname], Γ.namehint, newindex)
    (newname, Γ′)
end

"""
    split(Γ)

Separate `Γ` into two contexts containing the same free variables, but ensuring that different fresh
names are generated by them. 
"""
function split(Γ::NamingContext)
    l, r = splitindex(Γ.indices)
    NamingContext(Γ.freenames, Γ.namehint, l), NamingContext(Γ.freenames, Γ.namehint, r)
end
