Screen = {}
Screen.__index = Screen

local function round(n)
  return math.floor(n+0.5)
end

function Screen.create(bgColor)
  bgColor = bgColor or colors.black
  local screen = {}
  setmetatable(screen, Screen)

  screen.width,screen.height = term.getSize()
  for y=1,screen.height do
    screen[y] = {}
    for x=1,screen.width do
      screen[y][x] = bgColor
    end
  end
  return screen
end

function Screen:drawPixel(x,y,color)
  color = color or colors.white
  term.setCursorPos(x,y)
  term.setBackgroundColor(color)
  term.write(' ')
end

function Screen:setPixel(x,y,color)
  color = color or colors.white
  screen[y][x] = color
end

function Screen:draw()
  for y=1,self.height do
    term.setCursorPos(1,y)
    for x=1,self.width do
      if self[y][x] ~= 0 then
        term.setBackgroundColor(self[y][x])
        term.write(' ')
      else
        term.setCursorPos(x,y)
      end
    end
  end
end

function Screen:drawRectangle(xi,yi,xf,yf,color)
  local hLine = string.rep(' ',math.abs(xf-xi)+1)
  local leftX = math.min(xi,xf)
  term.setBackgroundColor(color)
  term.setCursorPos(leftX,yi)
  term.write(hLine)
  term.setCursorPos(leftX,yf)
  term.write(hLine)
  local minY = math.min(yi,yf)
  local maxY = math.max(yi,yf)
  for y=minY+1,maxY-1 do
    term.setCursorPos(xi,y)
    term.write(' ')
    term.setCursorPos(xf,y)
    term.write(' ')
  end
end

function Screen:setRectangle(xi,yi,xf,yf,color)
  local minX = math.min(xi,xf)
  local maxX = math.max(xi,xf)
  local minY = math.min(yi,yf)
  local maxY = math.max(yi,yf)
  for x=minX,maxX do
    self[yi][x] = color
  end
  for x=minX,maxX do
    self[yf][x] = color
  end
  for y=minY+1,maxY-1 do
    self[y][minX] = color
    self[y][maxX] = color
  end
end

function Screen:clearRectangle(xi,yi,xf,yf)
  local minX = math.min(xi,xf)
  local maxX = math.max(xi,xf)
  local minY = math.min(yi,yf)
  local maxY = math.max(yi,yf)
  term.setCursorPos(minX,yi)
  for x=minX,maxX do
    term.setBackgroundColor(self[yi][x])
    term.write(' ')
  end
  term.setCursorPos(minX,yf)
  for x=minX,maxX do
    term.setBackgroundColor(self[yf][x])
    term.write(' ')
  end
  for y=minY+1,maxY-1 do
    self:drawPixel(minX,y,self[y][minX])
    self:drawPixel(maxX,y,self[y][maxX])
  end
end

function Screen:drawCircle(xi,yi,xf,yf,color)
  color = color or colors.white
  local radius = round(math.sqrt((xf - xi)*(xf - xi) + (yf - yi)*(yf - yi)))

  for yt=-radius,radius do
    for xt=-radius,radius do
	  if (xi+xt >= 1 and xi+xt <= self.width) and (yi+yt >= 1 and yi+yt <= self.height) then
        if (radius == round(math.sqrt(xt*xt + yt*yt))) then
		  self:drawPixel(xi+xt,yi+yt,color)
		end
	  end
    end
  end
end

function Screen:setCircle(xi,yi,xf,yf,color)
  local radius = round(math.sqrt((xf - xi)*(xf - xi) + (yf - yi)*(yf - yi)))

  for yt=-radius,radius do
    for xt=-radius,radius do
	  if (xi+xt >= 1 and xi+xt <= self.width) and (yi+yt >= 1 and yi+yt <= self.height) then
        if (radius == round(math.sqrt(xt*xt + yt*yt))) then
		  self[yi+yt][xi+xt] = color
		end
	  end
    end
  end
end

function Screen:clearCircle(xi,yi,xf,yf)
  local radius = round(math.sqrt((xf - xi)*(xf - xi) + (yf - yi)*(yf - yi)))

  for yt=-radius,radius do
    for xt=-radius,radius do
	  if (xi+xt >= 1 and xi+xt <= self.width) and (yi+yt >= 1 and yi+yt <= self.height) then
        if (radius == round(math.sqrt(xt*xt + yt*yt))) then
		  self:drawPixel(xi+xt,yi+yt,self[yi+yt][xi+xt])
		end
	  end
    end
  end
end

local function dragFigure(screen, mode, xi, yi, color)
  local xf,yf = xi,yi
  while true do
    local e,p1,p2,p3 = os.pullEvent()
    if (e == "mouse_drag") then
      if mode == "rectangle" then
        screen:clearRectangle(xi,yi,xf,yf)
        xf,yf = p2,p3
        screen:drawRectangle(xi,yi,xf,yf,color)
	  elseif mode == "circle" then
	    screen:clearCircle(xi,yi,xf,yf)
		xf,yf = p2,p3
		screen:drawCircle(xi,yi,xf,yf,color)
      end
    elseif ((e == "mouse_click" and p1 == 1) or (e == "key" and p1 == keys.enter)) then
	  if mode == "rectangle" then
	    screen:setRectangle(xi,yi,xf,yf,color)
	  elseif mode == "circle" then
	    screen:setCircle(xi,yi,xf,yf,color)
	  end
      break
    elseif ((e == "mouse_click" and p1 == 2) or (e == "key" and p1 == keys.backspace)) then
      if mode == "rectangle" then
        screen:clearRectangle(xi,yi,xf,yf)
	  elseif mode == "circle" then
	    screen:clearCircle(xi,yi,xf,yf)
      end
      break
    end
  end
end

local s = Screen.create()
s:draw()
local mode = "circle"
local color = colors.blue

while true do
  local e,p1,p2,p3 = os.pullEvent('mouse_click')
  if (p1 == 1) then
    dragFigure(s, "rectangle", p2, p3, colors.red)
  elseif (p1 == 2) then
    dragFigure(s, "circle", p2, p3, colors.blue)
  end
end