using Oceananigans.Fields: compute_at!
using Oceananigans.OutputWriters: OutputWriters, _saveproperty!, serializeproperty!

# This is working just fine at the moment?
# But it will be veeeeery slow, as reconstruct_global_field is not
# a performant operation
function OutputWriters.fetch_output(mrf::MultiRegionField, model)
    field = reconstruct_global_field(mrf)
    compute_at!(field, model.clock.time)
    return parent(field)
end

OutputWriters.saveproperty!(file, address, p::Union{MultiRegionObject, MultiRegionField}) = _saveproperty!(file, address, p)

function OutputWriters.serializeproperty!(file, location, mrf::MultiRegionField{LX, LY, LZ}) where {LX, LY, LZ}
    p = reconstruct_global_field(mrf)

    serializeproperty!(file, location * "/location", (LX(), LY(), LZ()))
    serializeproperty!(file, location * "/data", parent(p))
    serializeproperty!(file, location * "/boundary_conditions", p.boundary_conditions)

    return nothing
end

function OutputWriters.serializeproperty!(file, location, mrg::MultiRegionGrids)
    file[location] = on_architecture(CPU(), reconstruct_global_grid(mrg))
end

#####
##### For a cubed sphere, we dump the entire field as is.
#####

function OutputWriters.fetch_output(csf::CubedSphereField, model)
    compute_at!(csf, model.clock.time)
    return parent(csf)
end

function OutputWriters.convert_output(mo::MultiRegionObject, writer)
    array_type = writer.array_type
    converted = Tuple(convert(array_type, obj) for obj in mo.regional_objects)
    return MultiRegionObject(converted)
end

function OutputWriters.construct_output(mrf::MultiRegionField{LX, LY, LZ}, user_indices, with_halos) where {LX, LY, LZ}
    multi_region_indices = output_indices(mrf, user_indices, with_halos)
    indices = getregion(multi_region_indices, 1)
    return Field(mrf; indices, NamedTuple()...)
end

function OutputWriters.serializeproperty!(file, location, csf::CubedSphereField{LX, LY, LZ}) where {LX, LY, LZ}
    csf_CPU = on_architecture(CPU(), csf)

    serializeproperty!(file, location * "/location", (LX(), LY(), LZ()))
    serializeproperty!(file, location * "/data", parent(csf_CPU))
    serializeproperty!(file, location * "/boundary_conditions", csf_CPU.boundary_conditions)

    return nothing
end

function OutputWriters.serializeproperty!(file, location, csg::ConformalCubedSphereGridOfSomeKind)
    file[location] = on_architecture(CPU(), csg)
end
