using Base64
using Serialization

OVERRIDE_MSG = "This should be overridden in implemention of AbstractEnvState"
num_rows(::Type{<:AbstractEnvState})              = @assert false OVERRIDE_MSG
num_cols(::Type{<:AbstractEnvState})              = @assert false OVERRIDE_MSG
player1_markers(::Type{<:AbstractEnvState})       = @assert false OVERRIDE_MSG
player2_markers(::Type{<:AbstractEnvState})       = @assert false OVERRIDE_MSG
get_pieces(::Type{<:AbstractEnvState})            = @assert false OVERRIDE_MSG
image_scale(::Type{<:AbstractEnvState})           = @assert false OVERRIDE_MSG
image_title_offset(::Type{<:AbstractEnvState})    = @assert false OVERRIDE_MSG
get_winner(::AbstractEnvState)                    = @assert false OVERRIDE_MSG
actions(::Type{<:AbstractEnvState})               = @assert false OVERRIDE_MSG
num_players(::Type{<:AbstractEnvState})           = @assert false OVERRIDE_MSG

key(s::AbstractEnvState) = encode(s)
svg(s::AbstractEnvState) = BoardStateSVG(s)
image_scale(x::AbstractEnvState) = image_scale(typeof(x))
image_title_offset(x::AbstractEnvState) = image_title_offset(typeof(x))
player1_markers(x::AbstractEnvState) = player1_markers(typeof(x))
player2_markers(x::AbstractEnvState) = player2_markers(typeof(x))
get_pieces(x::AbstractEnvState) = get_pieces(typeof(x))
num_rows(x::AbstractEnvState) = num_rows(typeof(x))
num_cols(x::AbstractEnvState) = num_cols(typeof(x))
num_cells(T::Type{<:AbstractEnvState}) = GridGames.getnumrows(T)*GridGames.getnumcols(T)
num_cells(x::AbstractEnvState) = num_cells(typeof(x))
Base.getindex(x::AbstractEnvState, idx::Integer) = x.board_state[idx]
Base.iterate(x::AbstractEnvState, state=1) = iterate(x.board_state, state)
action_space_size(T::Type{<:AbstractEnvState}) = getactions(T) |> length
action_space_size(x::AbstractEnvState) = action_space_size(typeof(x))
piece(s::AbstractEnvState, name::Symbol) = piece(typeof(s), name)
piece(::AbstractEnvState, ::Nothing) = nothing
piece(::Type{<:AbstractEnvState}, ::Nothing) = nothing
game_is_won(s::AbstractEnvState) = !isnothing(get_winner(s))
actions(s::AbstractEnvState) = actions(typeof(s))
to_play(s::AbstractEnvState) = s.to_play
dynamics(s::AbstractEnvState, a::Integer) = get_reward(s, a), next_state(s, a)
get_action_image(s::AbstractEnvState, a::Integer) = get_action_image(typeof(s), a)

function get_action_image(T::Type{<:AbstractEnvState}, a::Integer)
    channel = reshape(Flux.onehot(a, 1:num_cells(T)), GridGames.getnumrows(T), GridGames.getnumcols(T))
    Image(convert(Array{Float32}, Flux.unsqueeze(channel, 3)))
end

# function next_to_play(s::AbstractEnvState)
#     if num_players(typeof(s)) == 1
#         return PLAYER_1
#     end
#     to_play(s) == PLAYER_1 ? PLAYER_2 : PLAYER_1
# end

function get_image(s::AbstractEnvState)
    #pieces = piece_matrix(s)
    pieces = getpiecematrix(s)
    piece_channels = map(piece_idx->begin
        pieces .== piece_idx
    end, 1:length(getpieces(s)))
    to_play_multiplier = to_play(s).idx == PLAYER_1 ? Float32(1) : Float32(-1)
    to_play_channel = to_play_multiplier*ones(Float32, GridGames.getnumrows(s), GridGames.getnumcols(s))
    channels = reduce((a,b)->cat(a, b; dims=3), piece_channels; init=to_play_channel)
    Image(channels)
end
function get_reward(s::AbstractEnvState, action::Integer)
    winner = next_state(s, action) |> get_winner
    if isnothing(winner)
        return Float32(0)
    end
    winner == to_play(s) ? Float32(1) : Float32(-1)
end
function next_state(s::AbstractEnvState, action::Integer)
    typeof(s)(
        next_board_state(s, action),
        next_to_play(s)
    )
end
# function next_board_state(s::AbstractEnvState, action::Integer)
#     marker = get_marker(s, action)
#     coord = marker |> get_coordinate
#     byte_idx = coord_to_byte_idx(s, coord)
#     new_cell_byte = UInt8(get_piece_idx(s, action))
#     BoardState(s.board_state, byte_idx, new_cell_byte)
# end
get_marker(s::AbstractEnvState, a::Integer) = get_marker(s, to_play(s), a)
get_piece(s::AbstractEnvState, a::AbstractAction) = get_marker(s, a.i) |> get_piece


function get_piece_idx(s::AbstractEnvState, action::Integer)
    marker = get_marker(s, action)
    if isnothing(marker)
        return nothing
    end
    get_piece_idx(typeof(s), marker.piece.name)
end

function piece(T::Type{<:AbstractEnvState}, name::Symbol)
    get_pieces(T)[get_piece_idx(T, name)]
end

function get_piece(s::AbstractEnvState, name::Symbol)
    get_pieces(s)[get_piece_idx(s, name)]
end

function get_piece(s::AbstractEnvState, cell_idx::Integer)
    byte = s.board_state[cell_idx]
    if iszero(byte)
        return nothing
    end
    get_pieces(s)[byte]
end

function piece_count(s::AbstractEnvState, symbol)
    filter(byte->begin
        if iszero(byte)
            return false
        end
        get_pieces(s)[byte] == get_piece(s, symbol)
    end, s.board_state.bytes) |> length
end

get_piece_idx(s::AbstractEnvState, name::Symbol) = get_piece_idx(typeof(s), name)
function get_piece_idx(T::Type{<:AbstractEnvState}, name::Symbol)
    findfirst(x->x.name == name, get_pieces(T))
end

# @todo Remove these functions
hex_str(x::AbstractEnvState) = encode(x)[1:7]
#

# @todo Remove this. Needs to be at the game level to support 
# more flexible game play outcomes.
function last_reward(s::AbstractEnvState)
    isnothing(get_winner(s)) ? Float32(0) : Float32(1)
end

function Base.show(io::IO, s::AbstractEnvState)
    encoded = encode(s)
    if length(encoded) >= 40
        encoded = join([first(encoded, 20), "...", last(encoded, 20)])
    end
    print(io, "$(string(typeof(s))){\"$encoded\"}")
end

function Base.:-(s1::AbstractEnvState, s2::AbstractEnvState)
    bs1 = s1.board_state
    bs2 = s2.board_state
    byte_indices = findall(x->x > 0, bs1 - bs2)
    byte_idx_to_coord.(typeof(s1), byte_indices)
end

function byte_idx_to_coord(T::Type{<:AbstractEnvState}, byte_idx::Integer)
    row_idx = mod1(byte_idx, num_rows(T))
    col_idx = ((byte_idx - 1) ÷ num_rows(T)) + 1
    row_idx, col_idx
end

function coord_to_byte_idx(s::AbstractEnvState, coord::Tuple{<:Integer, <:Integer})
    coord_to_byte_idx(typeof(s), coord)
end

function coord_to_byte_idx(T::Type{<:AbstractEnvState}, coord::Tuple{<:Integer, <:Integer})
    row_idx, col_idx = coord
    (col_idx - 1)*num_rows(T) + row_idx
end

# function Base.getindex(s::AbstractEnvState, row_idx::Integer, col_idx::Integer)
#     byte_idx = coord_to_byte_idx(s, (row_idx, col_idx))
#     s[byte_idx]
# end

# function Base.getindex(s::AbstractEnvState, row_idx::Integer, ::Colon)
#     s[row_idx, 1:num_cols(s)]
# end

# function Base.getindex(s::AbstractEnvState, row_idx::Integer, col_range::UnitRange{<:Integer})
#     start_byte_idx = coord_to_byte_idx(s, (row_idx, col_range[1]))
#     end_byte_idx = coord_to_byte_idx(s, (row_idx, col_range[end]))
#     by = num_rows(s)
#     [s[idx] for idx in start_byte_idx:by:end_byte_idx]
# end

# function Base.getindex(s::AbstractEnvState, ::Colon, col_idx::Integer)
#     s[1:num_rows(s), col_idx]
# end

# @todo row_range should support step size (the intermediate value) as well
# function Base.getindex(s::AbstractEnvState, row_range::UnitRange{<:Integer}, col_idx::Integer)
#     start_byte_idx = coord_to_byte_idx(s, (row_range[1], col_idx))
#     end_byte_idx = coord_to_byte_idx(s, (row_range[end], col_idx))
#     [s[idx] for idx in start_byte_idx:end_byte_idx]
# end

# @todo 
# function Base.getindex(board_state::AbstractEnvState, ::Colon, col_idx::Integer)

#Base.copy(s::AbstractEnvState) = typeof(s)(copy(g.board_state))

function Base.show(io::IO, m::MIME"text/plain", s::AbstractEnvState)
    function cell_fn(_,_,row,col)
        piece_idx = s[row,col]
        piece_idx == 0 ? nothing : plaintext(get_pieces(s)[piece_idx])
    end
    g = Grid(num_cols(s), num_rows(s))
    show(io, "text/plain", g(cell_fn))
end

# function piece_matrix(s::AbstractEnvState)
#     reshape(
#         collect(s.board_state.bytes), 
#         num_rows(s), 
#         num_cols(s))
# end

function diff_coord(s1::AbstractEnvState, s2::AbstractEnvState)
    diff = (s1 - s2)
    isempty(diff) ? nothing : diff[1]
end

function targetsvg(input_state::AbstractEnvState, target_state::AbstractEnvState, target::Target)
    r = @sprintf "%.2f" reward(target)
    v = @sprintf "%.2f" value(target)
    highlight_coord = diff_coord(target_state, input_state)
    gridsvg(
        target_state;
        title="target",
        title_offset=(-210,-195),
        header="r: $r",
        header_color="#333",
        br_header="v: $v",
        br_header_color="#333",
        highlight=highlight_coord
    )
end

function gridsvg(
    s::AbstractEnvState;
    highlight=nothing,
    highlight_color="#aacef0",
    main_title="state",
    aspect_ratio=(num_cols(s), num_rows(s)),
    scale=image_scale(s),
    main_title_offset=image_title_offset(s),
    kwargs...)
    function cell_fn(_,_,row,col)
        piece_idx = s[row,col]
        if piece_idx == 0
            return nothing
        end
        p = get_pieces(s)[piece_idx]
        bg_color = highlight == (row,col) ? highlight_color : nothing
        (value=p.icon_symbol, color=p.icon_color, size=p.icon_size, bg_color=bg_color)
    end
    GridSVG(
        Grid(num_cols(s), num_rows(s), cell_fn);
        main_title=main_title,
        main_title_offset=main_title_offset,
        aspect_ratio=aspect_ratio,
        scale=scale,
        kwargs...
    )
end

function Base.show(io::IO, m::MIME"image/svg+xml", s::AbstractBoardStateSVG)
    show(io, m, gridsvg(s.x))
end