using Random
using Statistics

mutable struct ReplayBuffers
    writable_uuid::String
    d::Dict{String,ReplayBuffer}
end

function ReplayBuffers(writeable::ReplayBuffer)
    ReplayBuffers(
        string(writeable.uuid),
        Dict(
            string(writeable.uuid) => writeable   
        )
    )
end

function getgamestepquality(x::ReplayBuffers, total_steps)
    avg_step = getavgstep(x)
    if ismissing(avg_step) || iszero(total_steps)
        return missing
    end
    round(avg_step / total_steps, digits=2)
end


function calcnumreanalyzed(x::ReplayBuffers)
    addends = gettotalreanalyses.(values(x.d))
    isempty(addends) ? 0 : sum(addends)
end

function getavgstep(x::ReplayBuffers)
    steps = getfield.(values(x.d), :step_mean) 
    isempty(steps) ? 0 : round(mean(steps), digits=1)
end

function Base.length(x::ReplayBuffers)
    sum = 0
    for key in keys(x.d)
        sum += length(x.d[key])
    end
    sum
end

Base.lastindex(x::ReplayBuffers) = length(x)
Base.iterate(x::ReplayBuffers, state=1) = state > length(x) ? nothing : (x[state], state+1)
getbuffer(x::ReplayBuffers, uuid_str::String) = valornothing(x.d, uuid_str)
getwritablebuffer(x::ReplayBuffers) = getbuffer(x, x.writable_uuid)

function Base.getindex(x::ReplayBuffers, idx::Integer)
    for key in keys(x.d)
        buffer = x.d[key]
        len = length(buffer)
        if idx <= len
            return buffer[idx]
        end
        idx -= len
    end
    @assert false "This should not be accessible"
end

function Base.rand(x::ReplayBuffers)
    if isempty(x)
        return nothing
    end
    x[rand(1:length(x))]
end

function randtrainingbatch(x::ReplayBuffers, config::MuZeroConfig)
    while isempty(x)
    end
    batch_size = config.batch_size
    games = [rand(x) for _ in 1:batch_size]
    TrainingBatch(
        config.game_type,
        config.action_space_size, 
        [TrainingSample(config, g) for g in games])
end

function setbuffer!(x::ReplayBuffers, buffer::ReplayBuffer)
    x.d[string(buffer.uuid)] = buffer
    nothing
end

function setbuffer!(x::ReplayBuffers, buffer_path::String)
    try
        setbuffer!(x, deserializeexpanduser(buffer_path))
    catch e
        @error "setbuffer! errored out: ($(buffer_path))" exception=(e, catch_backtrace())
    end
end

hasbuffer(x::ReplayBuffers, uuid_str::String) = haskey(x.d, uuid_str)