Storage.jl

# Defines different storages for zarr arrays. Currently only regular files (DirectoryStore)
# and Dictionaries are supported

"""
    abstract type AbstractStore 

This the abstract supertype for all Zarr store implementations.  Currently only regular files ([`DirectoryStore`](@ref))
and Dictionaries are supported.

## Interface

All subtypes of `AbstractStore` must implement the following methods:

- [`storagesize(d::AbstractStore, p::AbstractString)`](@ref storagesize)
- [`subdirs(d::AbstractStore, p::AbstractString)`](@ref subdirs)
- [`subkeys(d::AbstractStore, p::AbstractString)`](@ref subkeys)
- [`isinitialized(d::AbstractStore, p::AbstractString)`](@ref isinitialized)
- [`storefromstring(::Type{<: AbstractStore}, s, _)`](@ref storefromstring)
- `Base.getindex(d::AbstractStore, i::AbstractString)`: return the data stored in key `i` as a Vector{UInt8}
- `Base.setindex!(d::AbstractStore, v, i::AbstractString)`: write the values in `v` to the key `i` of the given store `d`

They may optionally implement the following methods:

- [`store_read_strategy(s::AbstractStore)`](@ref store_read_strategy): return the read strategy for the given store.  See [`SequentialRead`](@ref) and [`ConcurrentRead`](@ref).
"""
abstract type AbstractStore end

# Define the interface

"""
    S3Store(bucket::String; aws=nothing)
 
An S3-backed Zarr store. Available after loading the `ZarrAWSS3Ext` extension.
"""
struct S3Store <: AbstractStore
    bucket::String
    aws::Any
end

function S3Store(args...)
    error("AWSS3 must be loaded to use S3Store. Try `using AWSS3`.")
end

"""
    storagesize(d::AbstractStore, p::AbstractString)

This function shall return the size of all data files in a store at path `p`.
"""
function storagesize end


"""
    Base.getindex(d::AbstractStore,i::String)

Returns the data stored in the given key as a Vector{UInt8}
"""
Base.getindex(d::AbstractStore,i::AbstractString) = error("getindex not implemented for store $(typeof(d))")

"""
    Base.setindex!(d::AbstractStore,v,i::String)

Writes the values in v to the given store and key.
"""
Base.setindex!(d::AbstractStore,v,i::AbstractString) = error("setindex not implemented for store $(typeof(d))")

"""
    subdirs(d::AbstractStore, p)

Returns a list of keys for children stores in the given store at path p.
"""
function subdirs end

"""
    subkeys(d::AbstractStore, p)

Returns the keys of files in the given store.
"""
function subkeys end 

# Function to construct the full path to a chunk given the base path, Cartesian Index i, and the chunk ecoding
store_readchunk(s::AbstractStore, p, i::CartesianIndex, e::AbstractChunkKeyEncoding) = s[p, citostring(e, i)]
store_deletechunk(s::AbstractStore, p, i::CartesianIndex, e::AbstractChunkKeyEncoding) = delete!(s, p, citostring(e, i))
store_writechunk(s::AbstractStore, v, p, i::CartesianIndex, e::AbstractChunkKeyEncoding) = s[p, citostring(e, i)] = v
store_isinitialized(s::AbstractStore, p, i::CartesianIndex, e::AbstractChunkKeyEncoding) = isinitialized(s, p, citostring(e, i))


#Functions to concat path and key 
Base.getindex(s::AbstractStore, p, i::AbstractString) = s[_concatpath(p, i)]
Base.delete!(s::AbstractStore, p, i::AbstractString) = delete!(s, _concatpath(p, i))
Base.haskey(s::AbstractStore, k::AbstractString) = isinitialized(s, k)
Base.setindex!(s::AbstractStore, v, p, i::AbstractString) = setindex!(s, v, _concatpath(p, i))


maybecopy(x) = copy(x)
maybecopy(x::String) = x


function getattrs(::ZarrFormat{2}, s::AbstractStore, p)
  atts = s[p,".zattrs"]
  if atts === nothing
    Dict()
  else
    JSON.parse(replace(String(maybecopy(atts)),": NaN,"=>": \"NaN\","); dicttype = Dict{String,Any})
  end
end

function getattrs(::ZarrFormat{3}, s::AbstractStore, p)
  md = s[p, "zarr.json"]
  if md === nothing
    error("zarr.json not found")
  else
    md = JSON.parse(replace(String(maybecopy(md)), ": NaN," => ": \"NaN\","); dicttype=Dict{String,Any})
    return get(md, "attributes", Dict{String,Any}())
  end
end

function writeattrs(::ZarrFormat{2}, s::AbstractStore, p, att::Dict; indent_json::Bool=false)
  b = IOBuffer()

  if indent_json
    JSON.print(b,att,4)
  else
    JSON.print(b,att)
  end

  s[p,".zattrs"] = take!(b)
  att
end

function writeattrs(::ZarrFormat{3}, s::AbstractStore, p, att::Dict; indent_json::Bool=false)
  # This is messy, we need to open zarr.json and replace the attributes section
  md = s[p, "zarr.json"]
  if md === nothing
    error("zarr.json not found")
  else
    md = JSON.parse(replace(String(maybecopy(md)), ": NaN," => ": \"NaN\","))
  end
  md = Dict(md)
  md["attributes"] = att

  b = IOBuffer()

  if indent_json
    JSON.print(b, md, 4)
  else
    JSON.print(b, md)
  end

  s[p, "zarr.json"] = take!(b)
  att
end

is_zarr3(s::AbstractStore, p) = isinitialized(s,_concatpath(p,"zarr.json"))
is_zarr2(s::AbstractStore, p) = is_zarray(ZarrFormat(Val(2)), s, p) || is_zgroup(ZarrFormat((Val(2))), s, p)

is_zgroup(::ZarrFormat{2}, s::AbstractStore, p) = isinitialized(s, _concatpath(p, ".zgroup"))
is_zarray(::ZarrFormat{2}, s::AbstractStore, p) = isinitialized(s, _concatpath(p, ".zarray"))
function is_zgroup(::ZarrFormat{3}, s::AbstractStore, p)
  isinitialized(s, _concatpath(p, "zarr.json")) || return false
  try
    metadata = getmetadata(ZarrFormat(Val(3)), s, p, false)
    metadata.node_type == "group"
  catch e
    e isa ArgumentError || rethrow()
    @warn "Skipping $p: $(e.msg)"
    false
  end
end
function is_zarray(::ZarrFormat{3}, s::AbstractStore, p)
  isinitialized(s, _concatpath(p, "zarr.json")) || return false
  try
    metadata = getmetadata(ZarrFormat(Val(3)), s, p, false)
    metadata.node_type == "array"
  catch e
    e isa ArgumentError || rethrow()
    @warn "Skipping $p: $(e.msg)"
    false
  end
end


isinitialized(s::AbstractStore, p, i::AbstractString) = isinitialized(s, _concatpath(p, i))
isinitialized(s::AbstractStore, i::AbstractString) = s[i] !== nothing

getmetadata(::ZarrFormat{2}, s::AbstractStore, p, fill_as_missing) = Metadata(String(maybecopy(s[p, ".zarray"])), fill_as_missing)

getmetadata(::ZarrFormat{3}, s::AbstractStore, p, fill_as_missing) = Metadata(String(maybecopy(s[p, "zarr.json"])), fill_as_missing)

function writemetadata(::ZarrFormat{2}, s::AbstractStore, p, m::AbstractMetadata; indent_json::Bool=false)
  met = IOBuffer()

  if indent_json
    JSON.print(met,m,4)
  else
    JSON.print(met,m)
  end
  
  s[p,".zarray"] = take!(met)
  m
end
function writemetadata(::ZarrFormat{3}, s::AbstractStore, p, m::AbstractMetadata; indent_json::Bool=false)
  met = IOBuffer()

  if indent_json
    JSON.print(met, m, 4)
  else
    JSON.print(met, m)
  end

  s[p, "zarr.json"] = take!(met)
  m
end



"""
    supports_partial_reads(::AbstractStore) -> Bool

Whether this store implements efficient byte-range reads via
[`read_range`](@ref) and [`getsize`](@ref). Defaults to `false`; stores
override to opt in (e.g. `DirectoryStore`).

The sharding partial-read fast path uses this to skip loading whole
shard files when only a few inner chunks are needed.
"""
supports_partial_reads(::AbstractStore) = false

"""
    read_range(s::AbstractStore, key::AbstractString, byte_range::UnitRange{Int})
        -> Union{Vector{UInt8}, Nothing}

Read just `byte_range` (1-based, inclusive) from the value stored under
`key`. Returns `nothing` if the key doesn't exist. Default implementation
falls back to `s[key][byte_range]`; stores that opt into
[`supports_partial_reads`](@ref) should override with a real partial read.
"""
function read_range(s::AbstractStore, key::AbstractString, byte_range::UnitRange{Int})
    bytes = s[key]
    bytes === nothing && return nothing
    return bytes[byte_range]
end

"""
    getsize(s::AbstractStore, key::AbstractString) -> Int

Size in bytes of the value stored under `key`, or 0 if it doesn't exist.
Default falls back to `length(s[key])`; partial-read-capable stores
should override with a cheap size lookup (e.g. `filesize`).
"""
function getsize(s::AbstractStore, key::AbstractString)
    bytes = s[key]
    bytes === nothing && return 0
    return length(bytes)
end


## Handling sequential vs parallel IO
struct SequentialRead end
struct ConcurrentRead
    ntasks::Int
end
store_read_strategy(::AbstractStore) = SequentialRead()

channelsize(s) = channelsize(store_read_strategy(s))
channelsize(::SequentialRead) = 0
channelsize(c::ConcurrentRead) = c.ntasks

read_items!(s::AbstractStore, c::AbstractChannel, e::AbstractChunkKeyEncoding, p, i) = read_items!(s, c, store_read_strategy(s), e, p, i)
function read_items!(s::AbstractStore, c::AbstractChannel, ::SequentialRead, e::AbstractChunkKeyEncoding, p, i)
    for ii in i
    res = store_readchunk(s, p, ii, e)
        put!(c,(ii=>res))
    end
end
function read_items!(s::AbstractStore, c::AbstractChannel, r::ConcurrentRead, e::AbstractChunkKeyEncoding, p, i)
    ntasks = r.ntasks
    #@show ntasks
    asyncmap(i,ntasks = ntasks) do ii
        #@show ii,objectid(current_task),p
    res = store_readchunk(s, p, ii, e)
        #@show ii,length(res)
        put!(c,(ii=>res))
        nothing
    end
end

write_items!(s::AbstractStore, c::AbstractChannel, e::AbstractChunkKeyEncoding, p, i) = write_items!(s, c, store_read_strategy(s), e, p, i)
function write_items!(s::AbstractStore, c::AbstractChannel, ::SequentialRead, e::AbstractChunkKeyEncoding, p, i)
  for _ in 1:length(i)
      ii,data = take!(c)
      if data === nothing
        if store_isinitialized(s, p, ii, e)
        store_deletechunk(s, p, ii, e)
        end
      else
      store_writechunk(s, data, p, ii, e)
      end
  end
  close(c)
end

function write_items!(s::AbstractStore, c::AbstractChannel, r::ConcurrentRead, e::AbstractChunkKeyEncoding, p, i)
  ntasks = r.ntasks
  asyncmap(i,ntasks = ntasks) do _
      ii,data = take!(c)
      if data === nothing
        if store_isinitialized(s, p, ii, e)
        store_deletechunk(s, p, ii, e)
        end
      else
      store_writechunk(s, data, p, ii, e) = data
      end
      nothing
  end
  close(c)
end

isemptysub(s::AbstractStore, p) = isempty(subkeys(s,p)) && isempty(subdirs(s,p))

#Here different storage backends can register regexes that are checked against
#during auto-check of storage format when doing zopen
storageregexlist = Pair[]
push!(storageregexlist, r"^s3://" => S3Store)

#include("formattedstore.jl")
include("directorystore.jl")
include("dictstore.jl")
include("gcstore.jl")
include("consolidated.jl")
include("http.jl")
include("zipstore.jl")