module IntCode

export Computer

import Base: getindex, setindex!

struct Memory
    data::Vector{Int}
end

function getindex(m::Memory, i::Int)
    i ≥ length(m.data) ? 0 : m.data[i + 1]
end

function setindex!(m::Memory, v::Int, i::Int)
    # expand memory by 50% of the requested value
    if i ≥ length(m.data)
        append!(m.data, zeros(round(Int, i * 1.5)))
    end
    setindex!(m.data, v, i + 1)
end

mutable struct Computer
    mem::Memory
    ptr::Int
    base::Int
    halt::Bool
end

Computer(program::Vector{Int}, ptr = 0, base = 0 , halt = false) =
    Computer(Memory(copy(program)), ptr, base, halt)

Computer(s::AbstractString) = Computer(parse.(Int, split(s, ',')))

function (c::Computer)(input::AbstractVector{<:Int})
    output = Int[]
    mem = c.mem

    function arg(idx)
        mode = mem[c.ptr] ÷ 10^(1 + idx) % 10
        # position mode
        if mode == 0
            return mem[c.ptr + idx]
        # immediate mode
        elseif mode == 1
            return c.ptr + idx
        # relative mode
        elseif mode == 2
            return mem[c.ptr + idx] + c.base
        else
            error("Mode $mode is not supported.")
        end
    end

    while c.halt == false
        opcode = mem[c.ptr] % 100

        # 1: addition
        if opcode == 1
            mem[arg(3)] = mem[arg(1)] + mem[arg(2)]
            c.ptr += 4
        # 2: multiplication
        elseif opcode == 2
            mem[arg(3)] = mem[arg(1)] * mem[arg(2)]
            c.ptr += 4
        # 3: input
        elseif opcode == 3
            if isempty(input)
                return output
            else
                mem[arg(1)] = popfirst!(input)
                c.ptr += 2
            end
        # 4: output
        elseif opcode == 4
            push!(output, mem[arg(1)])
            c.ptr += 2
        # 5: jump-if-true
        elseif opcode == 5
            if mem[arg(1)] != 0
                c.ptr = mem[arg(2)]
            else
                c.ptr += 3
            end
        # 6: jump-if-false
        elseif opcode == 6
            if mem[arg(1)] == 0
                c.ptr = mem[arg(2)]
            else
                c.ptr += 3
            end
        # 7: less than
        elseif opcode == 7
            mem[arg(3)] = mem[arg(1)] < mem[arg(2)] ? 1 : 0
            c.ptr += 4
        # 8: equals
        elseif opcode == 8
            mem[arg(3)] = mem[arg(1)] == mem[arg(2)] ? 1 : 0
            c.ptr += 4
        # 9: adjust relative relative base
        elseif opcode == 9
            c.base += mem[arg(1)]
            c.ptr += 2
        # 99: halt
        elseif opcode == 99
            c.halt = true
            return output
        else
            error("Program crashed: $opcode is not a valid opcode")
        end
    end
end

(c::Computer)(input...) = c(collect(input))

end